Skip to content

Commit 573e308

Browse files
arjunr2arjunr2
committed
Supported run CLI with async
1 parent 41af9bb commit 573e308

File tree

2 files changed

+134
-87
lines changed

2 files changed

+134
-87
lines changed

‎crates/wasmtime/src/runtime/rr/replay_driver.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,11 @@ impl<T: 'static> ReplayInstance<T> {
128128
&self.store
129129
}
130130

131+
/// Consume the [`ReplayInstance`] and extract the internal [`Store`]
132+
pub fn extract_store(self) -> Store<T> {
133+
self.store
134+
}
135+
131136
/// Run a single top-level event from the instance
132137
///
133138
/// "Top-level" events are those explicitly invoked events, namely:

‎src/commands/run.rs‎

Lines changed: 129 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ use std::io::BufReader;
1616
use std::path::{Path, PathBuf};
1717
use std::sync::{Arc, Mutex};
1818
use std::thread;
19+
#[cfg(feature = "rr")]
20+
use tokio::time::error::Elapsed;
1921
use wasi_common::sync::{Dir, TcpListener, WasiCtxBuilder, ambient_authority};
2022
#[cfg(feature = "rr")]
2123
use wasmtime::ReplayEnvironment;
@@ -101,16 +103,16 @@ impl RunCommand {
101103
) -> Result<()> {
102104
self.run.common.init_logging()?;
103105

104-
let mut config = self.run.common.config(None)?;
106+
#[cfg(feature = "rr")]
107+
let is_replaying = replay_opts.is_some();
105108
#[cfg(not(feature = "rr"))]
109+
let is_replaying = false;
110+
111+
let mut config = self.run.common.config(None)?;
106112
config.async_support(true);
107113
#[cfg(feature = "rr")]
108-
if replay_opts.is_some() {
109-
// Replay does not support async yet
110-
config.async_support(false);
114+
if is_replaying {
111115
config.rr(wasmtime::RRConfig::Replaying);
112-
} else {
113-
config.async_support(true);
114116
}
115117

116118
if self.run.common.wasm.timeout.is_some() {
@@ -173,13 +175,7 @@ impl RunCommand {
173175
};
174176

175177
let mut store = Store::new(&engine, host);
176-
#[cfg(not(feature = "rr"))]
177178
self.populate_with_wasi(&mut linker, &mut store, &main)?;
178-
// For replay, we don't populate WASI
179-
#[cfg(feature = "rr")]
180-
if !replay_opts.is_some() {
181-
self.populate_with_wasi(&mut linker, &mut store, &main)?;
182-
}
183179

184180
store.data_mut().limits = self.run.store_limits();
185181
store.limiter(|t| &mut t.limits);
@@ -192,32 +188,6 @@ impl RunCommand {
192188

193189
#[cfg(feature = "rr")]
194190
{
195-
// If this is a replay run, skip to replay setup and run to completion
196-
//
197-
// Note: Right now, replay doesn't inherit any store settings listed
198-
// above. This will have to change in the future. In general, replays will
199-
// need an "almost exact" superset of the run configurations, but with
200-
// certain different options (e.g. fuel consumption).
201-
if let Some(opts) = replay_opts {
202-
let settings = ReplaySettings {
203-
validate: opts.validate,
204-
deser_buffer_size: opts.deser_buffer_size,
205-
..Default::default()
206-
};
207-
208-
let mut renv = ReplayEnvironment::new(&engine, settings);
209-
match main {
210-
RunTarget::Core(m) => {
211-
renv.add_module(m);
212-
}
213-
RunTarget::Component(c) => {
214-
renv.add_component(c);
215-
}
216-
}
217-
let mut instance = renv.instantiate(BufReader::new(fs::File::open(opts.trace)?))?;
218-
return instance.run_to_completion();
219-
}
220-
221191
// Recording settings for this execution's store
222192
let record = &self.run.common.record;
223193
if let Some(path) = &record.path {
@@ -251,61 +221,133 @@ impl RunCommand {
251221
.wasm
252222
.timeout
253223
.unwrap_or(std::time::Duration::MAX);
254-
let result = runtime.block_on(async {
255-
tokio::time::timeout(dur, async {
256-
let mut profiled_modules: Vec<(String, Module)> = Vec::new();
257-
if let RunTarget::Core(m) = &main {
258-
profiled_modules.push(("".to_string(), m.clone()));
259-
}
260224

261-
// Load the preload wasm modules.
262-
for (name, path) in self.preloads.iter() {
263-
// Read the wasm module binary either as `*.wat` or a raw binary
264-
let preload_target = self.run.load_module(&engine, path)?;
265-
let preload_module = match preload_target {
266-
RunTarget::Core(m) => m,
267-
#[cfg(feature = "component-model")]
268-
RunTarget::Component(_) => {
269-
bail!("components cannot be loaded with `--preload`")
270-
}
271-
};
272-
profiled_modules.push((name.to_string(), preload_module.clone()));
273-
274-
// Add the module's functions to the linker.
275-
match &mut linker {
276-
#[cfg(feature = "cranelift")]
277-
CliLinker::Core(linker) => {
278-
linker
279-
.module_async(&mut store, name, &preload_module)
280-
.await
281-
.context(format!(
282-
"failed to process preload `{}` at `{}`",
283-
name,
284-
path.display()
285-
))?;
286-
}
287-
#[cfg(not(feature = "cranelift"))]
288-
CliLinker::Core(_) => {
289-
bail!("support for --preload disabled at compile time");
290-
}
291-
#[cfg(feature = "component-model")]
292-
CliLinker::Component(_) => {
293-
bail!("--preload cannot be used with components");
225+
let result = if is_replaying {
226+
// Note: While replaying, we need to retrieveIn general, replays will need an "almost exact" superset of
227+
// the run configurations, but with potentially certain different options (e.g. fuel consumption).
228+
#[cfg(feature = "rr")]
229+
{
230+
struct OkReplayT(Store<Host>);
231+
struct ErrReplayT(Error, Store<Host>);
232+
let opts = replay_opts.unwrap();
233+
let settings = ReplaySettings {
234+
validate: opts.validate,
235+
deser_buffer_size: opts.deser_buffer_size,
236+
..Default::default()
237+
};
238+
239+
let mut renv = ReplayEnvironment::new(&engine, settings);
240+
match &main {
241+
RunTarget::Core(m) => {
242+
renv.add_module(m.clone());
243+
}
244+
RunTarget::Component(c) => {
245+
renv.add_component(c.clone());
246+
}
247+
}
248+
let mut replay_instance = renv.instantiate_with_store(
249+
|| store,
250+
BufReader::new(fs::File::open(opts.trace)?),
251+
)?;
252+
253+
let result_replay_run: Result<Result<OkReplayT, ErrReplayT>, Elapsed> = runtime
254+
.block_on(async {
255+
tokio::time::timeout(dur, async {
256+
let res = replay_instance.run_to_completion_async().await;
257+
match res {
258+
Ok(_) => Ok(OkReplayT(replay_instance.extract_store())),
259+
Err(e) => Err(ErrReplayT(e, replay_instance.extract_store())),
260+
}
261+
})
262+
.await
263+
});
264+
match result_replay_run {
265+
Ok(Ok(OkReplayT(s))) => {
266+
store = s;
267+
Ok(Ok(()))
268+
}
269+
Ok(Err(ErrReplayT(e, s))) => {
270+
store = s;
271+
Ok(Err(e))
272+
}
273+
Err(elapsed) => {
274+
eprintln!(
275+
"Error: {:?}",
276+
Err::<(), Error>(anyhow::Error::from(wasmtime::Trap::Interrupt))
277+
.with_context(|| format!("timed out after {elapsed}"))
278+
);
279+
cfg_if::cfg_if! {
280+
if #[cfg(unix)] {
281+
std::process::exit(rustix::process::EXIT_SIGNALED_SIGABRT);
282+
} else if #[cfg(windows)] {
283+
// https://docs.microsoft.com/en-us/cpp/c-runtime-library/reference/abort?view=vs-2019
284+
std::process::exit(3);
285+
}
294286
}
295287
}
296288
}
289+
}
290+
#[cfg(not(feature = "rr"))]
291+
{
292+
unreachable!("should be unreachable when `rr` feature is disabled");
293+
}
294+
} else {
295+
runtime.block_on(async {
296+
tokio::time::timeout(dur, async {
297+
let mut profiled_modules: Vec<(String, Module)> = Vec::new();
298+
if let RunTarget::Core(m) = &main {
299+
profiled_modules.push(("".to_string(), m.clone()));
300+
}
297301

298-
self.load_main_module(&mut store, &mut linker, &main, profiled_modules)
299-
.await
300-
.with_context(|| {
301-
format!(
302-
"failed to run main module `{}`",
303-
self.module_and_args[0].to_string_lossy()
304-
)
305-
})
302+
// Load the preload wasm modules.
303+
for (name, path) in self.preloads.iter() {
304+
// Read the wasm module binary either as `*.wat` or a raw binary
305+
let preload_target = self.run.load_module(&engine, path)?;
306+
let preload_module = match preload_target {
307+
RunTarget::Core(m) => m,
308+
#[cfg(feature = "component-model")]
309+
RunTarget::Component(_) => {
310+
bail!("components cannot be loaded with `--preload`")
311+
}
312+
};
313+
profiled_modules.push((name.to_string(), preload_module.clone()));
314+
315+
// Add the module's functions to the linker.
316+
match &mut linker {
317+
#[cfg(feature = "cranelift")]
318+
CliLinker::Core(linker) => {
319+
linker
320+
.module_async(&mut store, name, &preload_module)
321+
.await
322+
.context(format!(
323+
"failed to process preload `{}` at `{}`",
324+
name,
325+
path.display()
326+
))?;
327+
}
328+
#[cfg(not(feature = "cranelift"))]
329+
CliLinker::Core(_) => {
330+
bail!("support for --preload disabled at compile time");
331+
}
332+
#[cfg(feature = "component-model")]
333+
CliLinker::Component(_) => {
334+
bail!("--preload cannot be used with components");
335+
}
336+
}
337+
}
338+
339+
self.load_main_module(&mut store, &mut linker, &main, profiled_modules)
340+
.await
341+
.with_context(|| {
342+
format!(
343+
"failed to run main module `{}`",
344+
self.module_and_args[0].to_string_lossy()
345+
)
346+
})
347+
})
348+
.await
306349
})
307-
.await
308-
});
350+
};
309351

310352
// Load the main wasm module.
311353
match result.unwrap_or_else(|elapsed| {

0 commit comments

Comments
 (0)