1use std::cell::RefCell;
8use std::future::Future;
9use std::io::{BufRead, BufReader, Error};
10use std::path::{Path, PathBuf};
11use std::pin::Pin;
12use std::process::Stdio;
13use std::rc::Rc;
14
15use bytes::{Bytes, BytesMut};
16use dfir_lang::graph::{AsCodeOptions, DfirGraph};
17use futures::{Sink, Stream};
18use serde::Serialize;
19use serde::de::DeserializeOwned;
20use stageleft::{QuotedWithContext, RuntimeData};
21
22use super::deploy_runtime_maelstrom::*;
23use crate::compile::builder::ExternalPortId;
24use crate::compile::deploy_provider::{ClusterSpec, Deploy, Node, RegisterPort};
25use crate::compile::trybuild::generate::{
26 ExampleBuildConfig, LinkingMode, TrybuildConfig, compile_trybuild_example,
27 create_graph_trybuild,
28};
29use crate::location::dynamic::LocationId;
30use crate::location::member_id::TaglessMemberId;
31use crate::location::{LocationKey, MembershipEvent, NetworkHint};
32
33pub enum MaelstromDeploy {}
40
41impl<'a> Deploy<'a> for MaelstromDeploy {
42 type Meta = ();
43 type InstantiateEnv = MaelstromDeployment;
44
45 type Process = MaelstromProcess;
46 type Cluster = MaelstromCluster;
47 type External = MaelstromExternal;
48
49 fn o2o_sink_source(
50 _env: &mut Self::InstantiateEnv,
51 _p1: &Self::Process,
52 _p1_port: &<Self::Process as Node>::Port,
53 _p2: &Self::Process,
54 _p2_port: &<Self::Process as Node>::Port,
55 _name: Option<&str>,
56 _networking_info: &crate::networking::NetworkingInfo,
57 _external_types: Option<(&syn::Type, &syn::Type)>,
58 ) -> (syn::Expr, syn::Expr) {
59 panic!("Maelstrom deployment does not support processes, only clusters")
60 }
61
62 fn o2o_connect(
63 _p1: &Self::Process,
64 _p1_port: &<Self::Process as Node>::Port,
65 _p2: &Self::Process,
66 _p2_port: &<Self::Process as Node>::Port,
67 ) -> Box<dyn FnOnce()> {
68 panic!("Maelstrom deployment does not support processes, only clusters")
69 }
70
71 fn o2m_sink_source(
72 _env: &mut Self::InstantiateEnv,
73 _p1: &Self::Process,
74 _p1_port: &<Self::Process as Node>::Port,
75 _c2: &Self::Cluster,
76 _c2_port: &<Self::Cluster as Node>::Port,
77 _name: Option<&str>,
78 _networking_info: &crate::networking::NetworkingInfo,
79 _external_types: Option<(&syn::Type, &syn::Type)>,
80 ) -> (syn::Expr, syn::Expr) {
81 panic!("Maelstrom deployment does not support processes, only clusters")
82 }
83
84 fn o2m_connect(
85 _p1: &Self::Process,
86 _p1_port: &<Self::Process as Node>::Port,
87 _c2: &Self::Cluster,
88 _c2_port: &<Self::Cluster as Node>::Port,
89 ) -> Box<dyn FnOnce()> {
90 panic!("Maelstrom deployment does not support processes, only clusters")
91 }
92
93 fn m2o_sink_source(
94 _env: &mut Self::InstantiateEnv,
95 _c1: &Self::Cluster,
96 _c1_port: &<Self::Cluster as Node>::Port,
97 _p2: &Self::Process,
98 _p2_port: &<Self::Process as Node>::Port,
99 _name: Option<&str>,
100 _networking_info: &crate::networking::NetworkingInfo,
101 _external_types: Option<(&syn::Type, &syn::Type)>,
102 ) -> (syn::Expr, syn::Expr) {
103 panic!("Maelstrom deployment does not support processes, only clusters")
104 }
105
106 fn m2o_connect(
107 _c1: &Self::Cluster,
108 _c1_port: &<Self::Cluster as Node>::Port,
109 _p2: &Self::Process,
110 _p2_port: &<Self::Process as Node>::Port,
111 ) -> Box<dyn FnOnce()> {
112 panic!("Maelstrom deployment does not support processes, only clusters")
113 }
114
115 fn m2m_sink_source(
116 env: &mut Self::InstantiateEnv,
117 _c1: &Self::Cluster,
118 _c1_port: &<Self::Cluster as Node>::Port,
119 _c2: &Self::Cluster,
120 _c2_port: &<Self::Cluster as Node>::Port,
121 _name: Option<&str>,
122 networking_info: &crate::networking::NetworkingInfo,
123 _external_types: Option<(&syn::Type, &syn::Type)>,
124 ) -> (syn::Expr, syn::Expr) {
125 use crate::networking::{NetworkingInfo, TcpFault};
126 match networking_info {
127 NetworkingInfo::Tcp { fault } => match (fault, env.nemesis.as_deref()) {
128 (TcpFault::Lossy | TcpFault::LossyDelayedForever, _) => {} (_, None) => {} (TcpFault::FailStop, Some("partition")) => {
131 panic!(
132 "Maelstrom partition nemesis requires lossy networking, but fail_stop was used. \
133 Use `TCP.lossy().bincode()` or `TCP.lossy_delayed_forever().bincode()` instead of `TCP.fail_stop().bincode()`."
134 );
135 }
136 (TcpFault::FailStop, Some(_)) => {} },
138 NetworkingInfo::Udp { .. } => {} }
140 deploy_maelstrom_m2m(RuntimeData::new("__hydro_lang_maelstrom_meta"))
141 }
142
143 fn m2m_connect(
144 _c1: &Self::Cluster,
145 _c1_port: &<Self::Cluster as Node>::Port,
146 _c2: &Self::Cluster,
147 _c2_port: &<Self::Cluster as Node>::Port,
148 ) -> Box<dyn FnOnce()> {
149 Box::new(|| {})
151 }
152
153 fn e2o_many_source(
154 _extra_stmts: &mut Vec<syn::Stmt>,
155 _p2: &Self::Process,
156 _p2_port: &<Self::Process as Node>::Port,
157 _codec_type: &syn::Type,
158 _shared_handle: String,
159 ) -> syn::Expr {
160 panic!("Maelstrom deployment does not support processes, only clusters")
161 }
162
163 fn e2o_many_sink(_shared_handle: String) -> syn::Expr {
164 panic!("Maelstrom deployment does not support processes, only clusters")
165 }
166
167 fn e2o_source(
168 _extra_stmts: &mut Vec<syn::Stmt>,
169 _p1: &Self::External,
170 _p1_port: &<Self::External as Node>::Port,
171 _p2: &Self::Process,
172 _p2_port: &<Self::Process as Node>::Port,
173 _codec_type: &syn::Type,
174 _shared_handle: String,
175 ) -> syn::Expr {
176 panic!("Maelstrom deployment does not support processes, only clusters")
177 }
178
179 fn e2o_connect(
180 _p1: &Self::External,
181 _p1_port: &<Self::External as Node>::Port,
182 _p2: &Self::Process,
183 _p2_port: &<Self::Process as Node>::Port,
184 _many: bool,
185 _server_hint: NetworkHint,
186 ) -> Box<dyn FnOnce()> {
187 panic!("Maelstrom deployment does not support processes, only clusters")
188 }
189
190 fn o2e_sink(
191 _p1: &Self::Process,
192 _p1_port: &<Self::Process as Node>::Port,
193 _p2: &Self::External,
194 _p2_port: &<Self::External as Node>::Port,
195 _shared_handle: String,
196 ) -> syn::Expr {
197 panic!("Maelstrom deployment does not support processes, only clusters")
198 }
199
200 fn cluster_ids(
201 _of_cluster: LocationKey,
202 ) -> impl QuotedWithContext<'a, &'a [TaglessMemberId], ()> + Clone + 'a {
203 cluster_members(RuntimeData::new("__hydro_lang_maelstrom_meta"), _of_cluster)
204 }
205
206 fn cluster_self_id() -> impl QuotedWithContext<'a, TaglessMemberId, ()> + Clone + 'a {
207 cluster_self_id(RuntimeData::new("__hydro_lang_maelstrom_meta"))
208 }
209
210 fn cluster_membership_stream(
211 _env: &mut Self::InstantiateEnv,
212 _at_location: &LocationId,
213 location_id: &LocationId,
214 ) -> impl QuotedWithContext<'a, Box<dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin>, ()>
215 {
216 cluster_membership_stream(location_id)
217 }
218}
219
220#[derive(Clone)]
222pub struct MaelstromProcess {
223 _private: (),
224}
225
226impl Node for MaelstromProcess {
227 type Port = String;
228 type Meta = ();
229 type InstantiateEnv = MaelstromDeployment;
230
231 fn next_port(&self) -> Self::Port {
232 panic!("Maelstrom deployment does not support processes")
233 }
234
235 fn update_meta(&self, _meta: &Self::Meta) {}
236
237 fn instantiate(
238 &self,
239 _env: &mut Self::InstantiateEnv,
240 _meta: &mut Self::Meta,
241 _graph: DfirGraph,
242 _extra_stmts: &[syn::Stmt],
243 _sidecars: &[syn::Expr],
244 _as_code_options: &AsCodeOptions,
245 ) {
246 panic!("Maelstrom deployment does not support processes")
247 }
248}
249
250#[derive(Clone)]
252pub struct MaelstromCluster {
253 next_port: Rc<RefCell<usize>>,
254 name_hint: Option<String>,
255}
256
257impl Node for MaelstromCluster {
258 type Port = String;
259 type Meta = ();
260 type InstantiateEnv = MaelstromDeployment;
261
262 fn next_port(&self) -> Self::Port {
263 let next_port = *self.next_port.borrow();
264 *self.next_port.borrow_mut() += 1;
265 format!("port_{}", next_port)
266 }
267
268 fn update_meta(&self, _meta: &Self::Meta) {}
269
270 fn instantiate(
271 &self,
272 env: &mut Self::InstantiateEnv,
273 _meta: &mut Self::Meta,
274 graph: DfirGraph,
275 extra_stmts: &[syn::Stmt],
276 sidecars: &[syn::Expr],
277 as_code_options: &AsCodeOptions,
278 ) {
279 let (bin_name, config) = create_graph_trybuild(
280 graph,
281 extra_stmts,
282 sidecars,
283 as_code_options,
284 self.name_hint.as_deref(),
285 crate::compile::trybuild::generate::DeployMode::Maelstrom,
286 LinkingMode::Dynamic,
287 );
288
289 env.bin_name = Some(bin_name);
290 env.trybuild = Some(config);
291 }
292}
293
294#[derive(Clone)]
296pub enum MaelstromExternal {}
297
298impl Node for MaelstromExternal {
299 type Port = String;
300 type Meta = ();
301 type InstantiateEnv = MaelstromDeployment;
302
303 fn next_port(&self) -> Self::Port {
304 unreachable!()
305 }
306
307 fn update_meta(&self, _meta: &Self::Meta) {}
308
309 fn instantiate(
310 &self,
311 _env: &mut Self::InstantiateEnv,
312 _meta: &mut Self::Meta,
313 _graph: DfirGraph,
314 _extra_stmts: &[syn::Stmt],
315 _sidecars: &[syn::Expr],
316 _as_code_options: &AsCodeOptions,
317 ) {
318 unreachable!()
319 }
320}
321
322impl<'a> RegisterPort<'a, MaelstromDeploy> for MaelstromExternal {
323 fn register(&self, _external_port_id: ExternalPortId, _port: Self::Port) {
324 unreachable!()
325 }
326
327 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
328 fn as_bytes_bidi(
329 &self,
330 _external_port_id: ExternalPortId,
331 ) -> impl Future<
332 Output = (
333 Pin<Box<dyn Stream<Item = Result<BytesMut, Error>>>>,
334 Pin<Box<dyn Sink<Bytes, Error = Error>>>,
335 ),
336 > + 'a {
337 async move { unreachable!() }
338 }
339
340 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
341 fn as_bincode_bidi<InT, OutT>(
342 &self,
343 _external_port_id: ExternalPortId,
344 ) -> impl Future<
345 Output = (
346 Pin<Box<dyn Stream<Item = OutT>>>,
347 Pin<Box<dyn Sink<InT, Error = Error>>>,
348 ),
349 > + 'a
350 where
351 InT: Serialize + 'static,
352 OutT: DeserializeOwned + 'static,
353 {
354 async move { unreachable!() }
355 }
356
357 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
358 fn as_bincode_sink<T: Serialize + 'static>(
359 &self,
360 _external_port_id: ExternalPortId,
361 ) -> impl Future<Output = Pin<Box<dyn Sink<T, Error = Error>>>> + 'a {
362 async move { unreachable!() }
363 }
364
365 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
366 fn as_bincode_source<T: DeserializeOwned + 'static>(
367 &self,
368 _external_port_id: ExternalPortId,
369 ) -> impl Future<Output = Pin<Box<dyn Stream<Item = T>>>> + 'a {
370 async move { unreachable!() }
371 }
372}
373
374#[derive(Clone)]
376pub struct MaelstromClusterSpec;
377
378impl<'a> ClusterSpec<'a, MaelstromDeploy> for MaelstromClusterSpec {
379 fn build(self, key: LocationKey, name_hint: &str) -> MaelstromCluster {
380 assert_eq!(
381 key,
382 LocationKey::FIRST,
383 "there should only be one location for a Maelstrom deployment"
384 );
385 MaelstromCluster {
386 next_port: Rc::new(RefCell::new(0)),
387 name_hint: Some(name_hint.to_owned()),
388 }
389 }
390}
391
392pub struct MaelstromDeployment {
397 pub node_count: usize,
399 pub maelstrom_path: PathBuf,
401 pub workload: String,
403 pub time_limit: Option<u64>,
405 pub rate: Option<u64>,
407 pub availability: Option<String>,
409 pub nemesis: Option<String>,
411 pub extra_args: Vec<String>,
413
414 pub(crate) bin_name: Option<String>,
416 pub(crate) trybuild: Option<TrybuildConfig>,
417}
418
419impl MaelstromDeployment {
420 pub fn new(workload: impl Into<String>) -> Self {
422 Self {
423 node_count: 1,
424 maelstrom_path: PathBuf::from("maelstrom"),
425 workload: workload.into(),
426 time_limit: None,
427 rate: None,
428 availability: None,
429 nemesis: None,
430 extra_args: vec![],
431 bin_name: None,
432 trybuild: None,
433 }
434 }
435
436 pub fn node_count(mut self, count: usize) -> Self {
438 self.node_count = count;
439 self
440 }
441
442 pub fn maelstrom_path(mut self, path: impl Into<PathBuf>) -> Self {
444 self.maelstrom_path = path.into();
445 self
446 }
447
448 pub fn time_limit(mut self, seconds: u64) -> Self {
450 self.time_limit = Some(seconds);
451 self
452 }
453
454 pub fn rate(mut self, rate: u64) -> Self {
456 self.rate = Some(rate);
457 self
458 }
459
460 pub fn availability(mut self, availability: impl Into<String>) -> Self {
462 self.availability = Some(availability.into());
463 self
464 }
465
466 pub fn nemesis(mut self, nemesis: impl Into<String>) -> Self {
468 self.nemesis = Some(nemesis.into());
469 self
470 }
471
472 pub fn extra_args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
474 self.extra_args.extend(args.into_iter().map(Into::into));
475 self
476 }
477
478 pub fn build(&self) -> Result<PathBuf, Error> {
486 let bin_name = self
487 .bin_name
488 .as_ref()
489 .expect("No binary name set - did you call deploy?");
490 let trybuild = self
491 .trybuild
492 .as_ref()
493 .expect("No trybuild config set - did you call deploy?");
494
495 let out = compile_trybuild_example(ExampleBuildConfig {
496 trybuild: trybuild.clone(),
497 bin_name: bin_name.clone(),
498 runtime_feature: "hydro___feature_maelstrom_runtime",
499 example_name: bin_name.clone(),
501 crate_type: None,
502 set_trybuild_lib_name: false,
503 allow_fuzz: false,
504 })
505 .map_err(|()| Error::other("Maelstrom binary compilation failed"))?;
506
507 out.keep().map_err(|e| Error::other(e.to_string()))
509 }
510
511 pub fn run(self) -> Result<(), Error> {
515 let binary_path = self.build()?;
516
517 std::process::Command::new(&binary_path)
526 .stdin(Stdio::null())
527 .stdout(Stdio::null())
528 .stderr(Stdio::null())
529 .spawn()?
530 .wait()?;
531
532 let run_dir = tempfile::tempdir().map_err(Error::other)?;
534
535 let mut cmd = std::process::Command::new(&self.maelstrom_path);
536 cmd.arg("test")
537 .arg("-w")
538 .arg(&self.workload)
539 .arg("--bin")
540 .arg(&binary_path)
541 .arg("--node-count")
542 .arg(self.node_count.to_string())
543 .current_dir(run_dir.path())
544 .stdout(Stdio::piped());
545
546 if let Some(time_limit) = self.time_limit {
547 cmd.arg("--time-limit").arg(time_limit.to_string());
548 }
549
550 if let Some(rate) = self.rate {
551 cmd.arg("--rate").arg(rate.to_string());
552 }
553
554 if let Some(availability) = self.availability {
555 cmd.arg("--availability").arg(availability);
556 }
557
558 if let Some(nemesis) = self.nemesis {
559 cmd.arg("--nemesis").arg(nemesis);
560 }
561
562 for arg in &self.extra_args {
563 cmd.arg(arg);
564 }
565
566 let spawned = cmd.spawn()?;
567
568 for line in BufReader::new(spawned.stdout.unwrap()).lines() {
569 let line = line?;
570 eprintln!("{}", line);
571
572 if line.starts_with("Analysis invalid!") {
573 let path = run_dir.keep();
574 dump_node_logs(&path);
575 return Err(Error::other(format!(
576 "Analysis was invalid. Maelstrom store at: {}",
577 path.display()
578 )));
579 } else if line.starts_with("Errors occurred during analysis, but no anomalies found.")
580 || line.starts_with("Everything looks good!")
581 {
582 return Ok(());
583 }
584 }
585
586 let path = run_dir.keep();
587 dump_node_logs(&path);
588 Err(Error::other(format!(
589 "Maelstrom produced an unexpected result. Store at: {}",
590 path.display()
591 )))
592 }
593
594 pub fn binary_path(&self) -> Option<PathBuf> {
596 self.build().ok()
597 }
598}
599
600fn dump_node_logs(run_dir: &Path) {
608 fn collect(dir: &Path, out: &mut Vec<PathBuf>) {
609 let Ok(entries) = std::fs::read_dir(dir) else {
610 return;
611 };
612 for entry in entries.flatten() {
613 let path = entry.path();
614 if path.is_symlink() || !path.is_dir() {
616 continue;
617 }
618 if path.file_name().is_some_and(|name| name == "node-logs") {
619 if let Ok(logs) = std::fs::read_dir(&path) {
620 out.extend(logs.flatten().map(|e| e.path()));
621 }
622 } else {
623 collect(&path, out);
624 }
625 }
626 }
627
628 let mut log_files = Vec::new();
629 collect(&run_dir.join("store"), &mut log_files);
630 log_files.sort();
631
632 for log in log_files {
633 eprintln!("==== Maelstrom node log: {} ====", log.display());
634 match std::fs::read_to_string(&log) {
635 Ok(contents) => eprint!("{}", contents),
636 Err(e) => eprintln!("(failed to read log: {})", e),
637 }
638 eprintln!("==== end of node log: {} ====", log.display());
639 }
640}