1 use std::io::{BufRead, Write};
2 use std::process::Command;
3 use wasmtime_wasi::cli::StdinStream;
4 use wasmtime_wasi::p2::Pollable;
5 
6 const VAR_NAME: &str = "__CHILD_PROCESS";
7 
main()8 fn main() {
9     if cfg!(miri) {
10         return;
11     }
12     // Skip this tests if it looks like we're in a cross-compiled situation and
13     // we're emulating this test for a different platform. In that scenario
14     // emulators (like QEMU) tend to not report signals the same way and such.
15     if wasmtime_test_util::cargo_test_runner().is_some() {
16         return;
17     }
18 
19     match std::env::var(VAR_NAME) {
20         Ok(_) => child_process(),
21         Err(_) => parent_process(),
22     }
23 
24     fn child_process() {
25         let mut result_write = std::io::stderr();
26         let mut child_running = true;
27         while child_running {
28             tokio::runtime::Builder::new_multi_thread()
29                 .enable_all()
30                 .build()
31                 .unwrap()
32                 .block_on(async {
33                     'task: loop {
34                         println!("child: creating stdin");
35                         let mut stdin = wasmtime_wasi::cli::stdin().p2_stream();
36 
37                         println!("child: checking that stdin is not ready");
38                         assert!(
39                             tokio::time::timeout(
40                                 std::time::Duration::from_millis(100),
41                                 stdin.ready()
42                             )
43                             .await
44                             .is_err(),
45                             "stdin available too soon"
46                         );
47 
48                         writeln!(&mut result_write, "start").unwrap();
49 
50                         println!("child: started");
51 
52                         let mut buffer = String::new();
53                         loop {
54                             println!("child: waiting for stdin to be ready");
55                             stdin.ready().await;
56 
57                             println!("child: reading input");
58                             // We can't effectively test for the case where stdin was closed, so panic if it is...
59                             let bytes = stdin.read(1024).unwrap();
60 
61                             println!("child got: {bytes:?}");
62 
63                             buffer.push_str(std::str::from_utf8(bytes.as_ref()).unwrap());
64                             if let Some((line, rest)) = buffer.split_once('\n') {
65                                 if line == "all done" {
66                                     writeln!(&mut result_write, "done").unwrap();
67                                     println!("child: exiting...");
68                                     child_running = false;
69                                     break 'task;
70                                 } else if line == "restart_runtime" {
71                                     writeln!(&mut result_write, "restarting").unwrap();
72                                     println!("child: restarting runtime...");
73                                     break 'task;
74                                 } else if line == "restart_task" {
75                                     writeln!(&mut result_write, "restarting").unwrap();
76                                     println!("child: restarting task...");
77                                     continue 'task;
78                                 } else {
79                                     writeln!(&mut result_write, "{line}").unwrap();
80                                 }
81 
82                                 buffer = rest.to_owned();
83                             }
84                         }
85                     }
86                 });
87             println!("child: runtime exited");
88         }
89         println!("child: exiting");
90     }
91 }
92 
parent_process()93 fn parent_process() {
94     let me = std::env::current_exe().unwrap();
95     let mut cmd = Command::new(me);
96     cmd.env(VAR_NAME, "1");
97     cmd.stdin(std::process::Stdio::piped());
98 
99     if std::env::args().any(|arg| arg == "--nocapture") {
100         cmd.stdout(std::process::Stdio::inherit());
101     } else {
102         cmd.stdout(std::process::Stdio::null());
103     }
104 
105     cmd.stderr(std::process::Stdio::piped());
106     let mut child = cmd.spawn().unwrap();
107 
108     let mut stdin_write = child.stdin.take().unwrap();
109     let mut result_read = std::io::BufReader::new(child.stderr.take().unwrap());
110 
111     let mut line = String::new();
112     result_read.read_line(&mut line).unwrap();
113     assert_eq!(line, "start\n");
114 
115     for i in 0..5 {
116         let message = format!("some bytes {i}\n");
117         stdin_write.write_all(message.as_bytes()).unwrap();
118         line.clear();
119         result_read.read_line(&mut line).unwrap();
120         assert_eq!(line, message);
121     }
122 
123     writeln!(&mut stdin_write, "restart_task").unwrap();
124     line.clear();
125     result_read.read_line(&mut line).unwrap();
126     assert_eq!(line, "restarting\n");
127     line.clear();
128 
129     result_read.read_line(&mut line).unwrap();
130     assert_eq!(line, "start\n");
131 
132     for i in 0..10 {
133         let message = format!("more bytes {i}\n");
134         stdin_write.write_all(message.as_bytes()).unwrap();
135         line.clear();
136         result_read.read_line(&mut line).unwrap();
137         assert_eq!(line, message);
138     }
139 
140     writeln!(&mut stdin_write, "restart_runtime").unwrap();
141     line.clear();
142     result_read.read_line(&mut line).unwrap();
143     assert_eq!(line, "restarting\n");
144     line.clear();
145 
146     result_read.read_line(&mut line).unwrap();
147     assert_eq!(line, "start\n");
148 
149     for i in 0..17 {
150         let message = format!("even more bytes {i}\n");
151         stdin_write.write_all(message.as_bytes()).unwrap();
152         line.clear();
153         result_read.read_line(&mut line).unwrap();
154         assert_eq!(line, message);
155     }
156 
157     writeln!(&mut stdin_write, "all done").unwrap();
158 
159     line.clear();
160     result_read.read_line(&mut line).unwrap();
161     assert_eq!(line, "done\n");
162 }
163