deterministic_wasi_ctx/
scheduling.rs

1use std::{mem, ptr, slice};
2
3use crate::wasi_abi::{
4    Event, EventFdReadwrite, Subscription, ERRNO_SUCCESS, EVENTTYPE_CLOCK, EVENTTYPE_FD_READ,
5    EVENTTYPE_FD_WRITE,
6};
7use wasmtime::{format_err, Caller, Linker, Result};
8
9/// Adds implementations for WASI preview 1 `poll_oneoff` and `sched_yield` to
10/// the linker which will return immediately.
11/// Note: This function will enable shadowing on the linker.
12pub fn replace_scheduling_functions<T>(linker: &mut Linker<T>) -> Result<()>
13where
14    T: Send + 'static,
15{
16    override_scheduling_functions(linker, "wasi_snapshot_preview1")
17}
18
19/// Adds implementations for WASI preview 0 `poll_oneoff` and `sched_yield` to
20/// the linker which will return immediately.
21/// Note: This function will enable shadowing on the linker.
22pub fn replace_scheduling_functions_for_wasi_preview_0<T>(linker: &mut Linker<T>) -> Result<()>
23where
24    T: Send + 'static,
25{
26    override_scheduling_functions(linker, "wasi_unstable")
27}
28
29fn override_scheduling_functions<T: 'static>(linker: &mut Linker<T>, module: &str) -> Result<()> {
30    linker.allow_shadowing(true);
31    linker.func_wrap(
32        module,
33        "poll_oneoff",
34        |mut caller: Caller<'_, T>,
35         in_ptr: i32,
36         out_ptr: i32,
37         nsubscriptions: i32,
38         nevents_ptr: i32|
39         -> Result<i32> {
40            let in_ptr = in_ptr as usize;
41            let out_ptr = out_ptr as usize;
42            let nsubscriptions = nsubscriptions as usize;
43            let nevents_ptr = nevents_ptr as usize;
44            // See https://github.com/WebAssembly/WASI/blob/3d5e0553cd01dd4d6e2c06ad2a702ee9dda17b7f/legacy/tools/witx-docs.md#pointers
45            let memory = caller
46                .get_export("memory")
47                .map_or_else(|| Err(format_err!("missing required memory export")), Ok)?
48                .into_memory()
49                .map_or_else(|| Err(format_err!("missing required memory export")), Ok)?;
50
51            for i in 0..nsubscriptions {
52                // Read the `Subscription` from memory from `in_ptr`.
53                let offset = in_ptr + (i * mem::size_of::<Subscription>());
54                let mut subscription_buffer = [0u8; mem::size_of::<Subscription>()];
55                memory.read(&caller, offset, &mut subscription_buffer)?;
56                let subscription =
57                    unsafe { ptr::read(subscription_buffer.as_ptr() as *const Subscription) };
58
59                // Create a successful `Event` for each subscription.
60                let event = Event {
61                    userdata: subscription.userdata,
62                    error: ERRNO_SUCCESS,
63                    // See https://github.com/WebAssembly/wasi-libc/blob/e9524a0980b9bb6bb92e87a41ed1055bdda5bb86/libc-bottom-half/headers/public/wasi/api.h#L1100-L1121
64                    // for the mapping between the integers and the event type.
65                    type_: match subscription.u.tag {
66                        0 => EVENTTYPE_CLOCK,
67                        1 => EVENTTYPE_FD_READ,
68                        2 => EVENTTYPE_FD_WRITE,
69                        _ => unreachable!(),
70                    },
71                    fd_readwrite: EventFdReadwrite {
72                        nbytes: 0,
73                        flags: 0,
74                    },
75                };
76
77                // Write the event into memory at `out_ptr`.
78                let offset = out_ptr + (i * mem::size_of::<Event>());
79                let event_buffer = unsafe {
80                    slice::from_raw_parts(
81                        &event as *const Event as *const u8,
82                        mem::size_of::<Event>(),
83                    )
84                };
85                memory.write(&mut caller, offset, event_buffer)?
86            }
87
88            // Copy number of subscriptions into number of events pointer.
89            let buffer = nsubscriptions.to_le_bytes();
90            memory.write(&mut caller, nevents_ptr, &buffer)?;
91
92            Ok(ERRNO_SUCCESS as i32)
93        },
94    )?;
95
96    linker.func_wrap(module, "sched_yield", || ERRNO_SUCCESS as i32)?;
97
98    Ok(())
99}
deterministic_wasi_ctx/scheduling.rs

deterministic_wasi_ctx/
scheduling.rs
