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::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 ) {
245 panic!("Maelstrom deployment does not support processes")
246 }
247}
248
249#[derive(Clone)]
251pub struct MaelstromCluster {
252 next_port: Rc<RefCell<usize>>,
253 name_hint: Option<String>,
254}
255
256impl Node for MaelstromCluster {
257 type Port = String;
258 type Meta = ();
259 type InstantiateEnv = MaelstromDeployment;
260
261 fn next_port(&self) -> Self::Port {
262 let next_port = *self.next_port.borrow();
263 *self.next_port.borrow_mut() += 1;
264 format!("port_{}", next_port)
265 }
266
267 fn update_meta(&self, _meta: &Self::Meta) {}
268
269 fn instantiate(
270 &self,
271 env: &mut Self::InstantiateEnv,
272 _meta: &mut Self::Meta,
273 graph: DfirGraph,
274 extra_stmts: &[syn::Stmt],
275 sidecars: &[syn::Expr],
276 ) {
277 let (bin_name, config) = create_graph_trybuild(
278 graph,
279 extra_stmts,
280 sidecars,
281 self.name_hint.as_deref(),
282 crate::compile::trybuild::generate::DeployMode::Maelstrom,
283 LinkingMode::Dynamic,
284 );
285
286 env.bin_name = Some(bin_name);
287 env.trybuild = Some(config);
288 }
289}
290
291#[derive(Clone)]
293pub enum MaelstromExternal {}
294
295impl Node for MaelstromExternal {
296 type Port = String;
297 type Meta = ();
298 type InstantiateEnv = MaelstromDeployment;
299
300 fn next_port(&self) -> Self::Port {
301 unreachable!()
302 }
303
304 fn update_meta(&self, _meta: &Self::Meta) {}
305
306 fn instantiate(
307 &self,
308 _env: &mut Self::InstantiateEnv,
309 _meta: &mut Self::Meta,
310 _graph: DfirGraph,
311 _extra_stmts: &[syn::Stmt],
312 _sidecars: &[syn::Expr],
313 ) {
314 unreachable!()
315 }
316}
317
318impl<'a> RegisterPort<'a, MaelstromDeploy> for MaelstromExternal {
319 fn register(&self, _external_port_id: ExternalPortId, _port: Self::Port) {
320 unreachable!()
321 }
322
323 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
324 fn as_bytes_bidi(
325 &self,
326 _external_port_id: ExternalPortId,
327 ) -> impl Future<
328 Output = (
329 Pin<Box<dyn Stream<Item = Result<BytesMut, Error>>>>,
330 Pin<Box<dyn Sink<Bytes, Error = Error>>>,
331 ),
332 > + 'a {
333 async move { unreachable!() }
334 }
335
336 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
337 fn as_bincode_bidi<InT, OutT>(
338 &self,
339 _external_port_id: ExternalPortId,
340 ) -> impl Future<
341 Output = (
342 Pin<Box<dyn Stream<Item = OutT>>>,
343 Pin<Box<dyn Sink<InT, Error = Error>>>,
344 ),
345 > + 'a
346 where
347 InT: Serialize + 'static,
348 OutT: DeserializeOwned + 'static,
349 {
350 async move { unreachable!() }
351 }
352
353 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
354 fn as_bincode_sink<T: Serialize + 'static>(
355 &self,
356 _external_port_id: ExternalPortId,
357 ) -> impl Future<Output = Pin<Box<dyn Sink<T, Error = Error>>>> + 'a {
358 async move { unreachable!() }
359 }
360
361 #[expect(clippy::manual_async_fn, reason = "false positive, involves lifetimes")]
362 fn as_bincode_source<T: DeserializeOwned + 'static>(
363 &self,
364 _external_port_id: ExternalPortId,
365 ) -> impl Future<Output = Pin<Box<dyn Stream<Item = T>>>> + 'a {
366 async move { unreachable!() }
367 }
368}
369
370#[derive(Clone)]
372pub struct MaelstromClusterSpec;
373
374impl<'a> ClusterSpec<'a, MaelstromDeploy> for MaelstromClusterSpec {
375 fn build(self, key: LocationKey, name_hint: &str) -> MaelstromCluster {
376 assert_eq!(
377 key,
378 LocationKey::FIRST,
379 "there should only be one location for a Maelstrom deployment"
380 );
381 MaelstromCluster {
382 next_port: Rc::new(RefCell::new(0)),
383 name_hint: Some(name_hint.to_owned()),
384 }
385 }
386}
387
388pub struct MaelstromDeployment {
393 pub node_count: usize,
395 pub maelstrom_path: PathBuf,
397 pub workload: String,
399 pub time_limit: Option<u64>,
401 pub rate: Option<u64>,
403 pub availability: Option<String>,
405 pub nemesis: Option<String>,
407 pub extra_args: Vec<String>,
409
410 pub(crate) bin_name: Option<String>,
412 pub(crate) trybuild: Option<TrybuildConfig>,
413}
414
415impl MaelstromDeployment {
416 pub fn new(workload: impl Into<String>) -> Self {
418 Self {
419 node_count: 1,
420 maelstrom_path: PathBuf::from("maelstrom"),
421 workload: workload.into(),
422 time_limit: None,
423 rate: None,
424 availability: None,
425 nemesis: None,
426 extra_args: vec![],
427 bin_name: None,
428 trybuild: None,
429 }
430 }
431
432 pub fn node_count(mut self, count: usize) -> Self {
434 self.node_count = count;
435 self
436 }
437
438 pub fn maelstrom_path(mut self, path: impl Into<PathBuf>) -> Self {
440 self.maelstrom_path = path.into();
441 self
442 }
443
444 pub fn time_limit(mut self, seconds: u64) -> Self {
446 self.time_limit = Some(seconds);
447 self
448 }
449
450 pub fn rate(mut self, rate: u64) -> Self {
452 self.rate = Some(rate);
453 self
454 }
455
456 pub fn availability(mut self, availability: impl Into<String>) -> Self {
458 self.availability = Some(availability.into());
459 self
460 }
461
462 pub fn nemesis(mut self, nemesis: impl Into<String>) -> Self {
464 self.nemesis = Some(nemesis.into());
465 self
466 }
467
468 pub fn extra_args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
470 self.extra_args.extend(args.into_iter().map(Into::into));
471 self
472 }
473
474 pub fn build(&self) -> Result<PathBuf, Error> {
482 let bin_name = self
483 .bin_name
484 .as_ref()
485 .expect("No binary name set - did you call deploy?");
486 let trybuild = self
487 .trybuild
488 .as_ref()
489 .expect("No trybuild config set - did you call deploy?");
490
491 let out = compile_trybuild_example(ExampleBuildConfig {
492 trybuild: trybuild.clone(),
493 bin_name: bin_name.clone(),
494 runtime_feature: "hydro___feature_maelstrom_runtime",
495 example_name: bin_name.clone(),
497 crate_type: None,
498 set_trybuild_lib_name: false,
499 allow_fuzz: false,
500 })
501 .map_err(|()| Error::other("Maelstrom binary compilation failed"))?;
502
503 out.keep().map_err(|e| Error::other(e.to_string()))
505 }
506
507 pub fn run(self) -> Result<(), Error> {
511 let binary_path = self.build()?;
512
513 std::process::Command::new(&binary_path)
522 .stdin(Stdio::null())
523 .stdout(Stdio::null())
524 .stderr(Stdio::null())
525 .spawn()?
526 .wait()?;
527
528 let run_dir = tempfile::tempdir().map_err(Error::other)?;
530
531 let mut cmd = std::process::Command::new(&self.maelstrom_path);
532 cmd.arg("test")
533 .arg("-w")
534 .arg(&self.workload)
535 .arg("--bin")
536 .arg(&binary_path)
537 .arg("--node-count")
538 .arg(self.node_count.to_string())
539 .current_dir(run_dir.path())
540 .stdout(Stdio::piped());
541
542 if let Some(time_limit) = self.time_limit {
543 cmd.arg("--time-limit").arg(time_limit.to_string());
544 }
545
546 if let Some(rate) = self.rate {
547 cmd.arg("--rate").arg(rate.to_string());
548 }
549
550 if let Some(availability) = self.availability {
551 cmd.arg("--availability").arg(availability);
552 }
553
554 if let Some(nemesis) = self.nemesis {
555 cmd.arg("--nemesis").arg(nemesis);
556 }
557
558 for arg in &self.extra_args {
559 cmd.arg(arg);
560 }
561
562 let spawned = cmd.spawn()?;
563
564 for line in BufReader::new(spawned.stdout.unwrap()).lines() {
565 let line = line?;
566 eprintln!("{}", &line);
567
568 if line.starts_with("Analysis invalid!") {
569 let path = run_dir.keep();
570 dump_node_logs(&path);
571 return Err(Error::other(format!(
572 "Analysis was invalid. Maelstrom store at: {}",
573 path.display()
574 )));
575 } else if line.starts_with("Errors occurred during analysis, but no anomalies found.")
576 || line.starts_with("Everything looks good!")
577 {
578 return Ok(());
579 }
580 }
581
582 let path = run_dir.keep();
583 dump_node_logs(&path);
584 Err(Error::other(format!(
585 "Maelstrom produced an unexpected result. Store at: {}",
586 path.display()
587 )))
588 }
589
590 pub fn binary_path(&self) -> Option<PathBuf> {
592 self.build().ok()
593 }
594}
595
596fn dump_node_logs(run_dir: &Path) {
604 fn collect(dir: &Path, out: &mut Vec<PathBuf>) {
605 let Ok(entries) = std::fs::read_dir(dir) else {
606 return;
607 };
608 for entry in entries.flatten() {
609 let path = entry.path();
610 if path.is_symlink() || !path.is_dir() {
612 continue;
613 }
614 if path.file_name().is_some_and(|name| name == "node-logs") {
615 if let Ok(logs) = std::fs::read_dir(&path) {
616 out.extend(logs.flatten().map(|e| e.path()));
617 }
618 } else {
619 collect(&path, out);
620 }
621 }
622 }
623
624 let mut log_files = Vec::new();
625 collect(&run_dir.join("store"), &mut log_files);
626 log_files.sort();
627
628 for log in log_files {
629 eprintln!("==== Maelstrom node log: {} ====", log.display());
630 match std::fs::read_to_string(&log) {
631 Ok(contents) => eprint!("{}", contents),
632 Err(e) => eprintln!("(failed to read log: {})", e),
633 }
634 eprintln!("==== end of node log: {} ====", log.display());
635 }
636}