async_wasi/snapshots/preview_1/
async_poll.rs

1use crate::snapshots::{
2    common::{
3        memory::{Memory, WasmPtr},
4        net::{self, ConnectState, SubscriptionClock},
5        types::*,
6    },
7    Errno, WasiCtx,
8};
9use futures::{stream::FuturesUnordered, StreamExt};
10use net::{async_tokio::AsyncWasiSocket, PrePoll, SubscriptionFd, SubscriptionFdType};
11use std::time::Duration;
12
13fn handle_event_err(type_: SubscriptionFdType, errno: Errno) -> __wasi_event_t {
14    let mut r = __wasi_event_t {
15        userdata: 0,
16        error: errno.0,
17        type_: 0,
18        fd_readwrite: __wasi_event_fd_readwrite_t {
19            nbytes: 0,
20            flags: 0,
21        },
22    };
23    match type_ {
24        SubscriptionFdType::Read(userdata) => {
25            r.userdata = userdata;
26            r.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_FD_READ;
27        }
28        SubscriptionFdType::Write(userdata) => {
29            r.userdata = userdata;
30            r.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_FD_WRITE;
31        }
32        SubscriptionFdType::Both { read: userdata, .. } => {
33            r.userdata = userdata;
34            r.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_FD_READ;
35        }
36    }
37    r
38}
39
40async fn wait_fd(
41    fd_index: usize,
42    socket: &AsyncWasiSocket,
43    type_: SubscriptionFdType,
44) -> Result<(__wasi_event_t, Option<usize>), Errno> {
45    let connecting = ConnectState::Connecting == socket.state.so_conn_state;
46
47    let handler = |r: Result<(), std::io::Error>, userdata, type_| {
48        log::trace!("wait_fd {fd_index} {r:?}");
49        match r {
50            Ok(_) => (
51                __wasi_event_t {
52                    userdata,
53                    error: 0,
54                    type_,
55                    fd_readwrite: __wasi_event_fd_readwrite_t {
56                        nbytes: 0,
57                        flags: 0,
58                    },
59                },
60                if connecting { Some(fd_index) } else { None },
61            ),
62            Err(e) => (
63                __wasi_event_t {
64                    userdata,
65                    error: Errno::from(e).0,
66                    type_,
67                    fd_readwrite: __wasi_event_fd_readwrite_t {
68                        nbytes: 0,
69                        flags: __wasi_eventrwflags_t::__WASI_EVENTRWFLAGS_FD_READWRITE_HANGUP,
70                    },
71                },
72                None,
73            ),
74        }
75    };
76
77    match type_ {
78        SubscriptionFdType::Write(userdata) => {
79            let write_result = socket.writable().await;
80            log::trace!("wait_fd {fd_index} writeable");
81
82            Ok(handler(
83                write_result,
84                userdata,
85                __wasi_eventtype_t::__WASI_EVENTTYPE_FD_WRITE,
86            ))
87        }
88        SubscriptionFdType::Read(userdata) => {
89            let read_result = socket.readable().await;
90            log::trace!("wait_fd {fd_index} readable");
91
92            Ok(handler(
93                read_result,
94                userdata,
95                __wasi_eventtype_t::__WASI_EVENTTYPE_FD_READ,
96            ))
97        }
98        SubscriptionFdType::Both { read, write } => {
99            tokio::select! {
100                read_result=socket.readable()=>{
101                    log::trace!("wait_fd {fd_index} readable");
102
103                    Ok(handler(
104                        read_result,
105                        read,
106                        __wasi_eventtype_t::__WASI_EVENTTYPE_FD_READ,
107                    ))
108                }
109                write_result=socket.writable()=>{
110                    log::trace!("wait_fd {fd_index} writeable");
111
112                    Ok(handler(
113                        write_result,
114                        write,
115                        __wasi_eventtype_t::__WASI_EVENTTYPE_FD_WRITE,
116                    ))
117                }
118            }
119        }
120    }
121}
122
123async fn poll_only_fd<M: Memory>(
124    ctx: &mut WasiCtx,
125    mem: &mut M,
126    out_ptr: WasmPtr<__wasi_event_t>,
127    nsubscriptions: usize,
128    revents_num_ptr: WasmPtr<__wasi_size_t>,
129    fd_vec: Vec<SubscriptionFd>,
130) -> Result<(), Errno> {
131    if fd_vec.is_empty() {
132        mem.write_data(revents_num_ptr, 0)?;
133    } else {
134        let r_events = mem.mut_slice(out_ptr, nsubscriptions)?;
135        let mut wait = FuturesUnordered::new();
136        let mut i = 0;
137        for SubscriptionFd { fd, type_ } in fd_vec {
138            match ctx.vfs.get_socket(fd as usize) {
139                Ok(s) => {
140                    wait.push(wait_fd(fd as usize, s, type_));
141                }
142                Err(e) => {
143                    r_events[i] = handle_event_err(type_, e);
144                    i += 1;
145                }
146            }
147        }
148
149        if i == 0 {
150            let mut connected_fds = vec![];
151
152            let (v, connected_fd) = wait.select_next_some().await?;
153            connected_fds.push(connected_fd);
154            r_events[i] = v;
155            i += 1;
156
157            'wait_poll: loop {
158                if i >= nsubscriptions {
159                    break 'wait_poll;
160                }
161                futures::select! {
162                    v = wait.next() => {
163                        if let Some(v) = v {
164                            let (v, connected_fd) = v?;
165                            connected_fds.push(connected_fd);
166                            r_events[i] = v;
167                            i += 1;
168                        } else {
169                            break 'wait_poll;
170                        }
171                    }
172                    default => {
173                        break 'wait_poll;
174                    }
175                };
176            }
177
178            drop(wait);
179
180            for fd in connected_fds.into_iter().flatten() {
181                if let Ok(socket) = ctx.vfs.get_mut_socket(fd) {
182                    socket.state.so_conn_state = ConnectState::Connected;
183                    socket.writable.set_writable();
184                }
185            }
186        }
187
188        mem.write_data(revents_num_ptr, i as u32)?;
189    }
190    Ok(())
191}
192
193async fn poll_fd_timeout<M: Memory>(
194    ctx: &mut WasiCtx,
195    mem: &mut M,
196    out_ptr: WasmPtr<__wasi_event_t>,
197    nsubscriptions: usize,
198    revents_num_ptr: WasmPtr<__wasi_size_t>,
199    clock: SubscriptionClock,
200    fd_vec: Vec<SubscriptionFd>,
201) -> Result<(), Errno> {
202    let r_events = mem.mut_slice(out_ptr, nsubscriptions)?;
203    let mut wait = FuturesUnordered::new();
204
205    let mut i = 0;
206
207    for SubscriptionFd { fd, type_ } in fd_vec {
208        match ctx.vfs.get_socket(fd as usize) {
209            Ok(s) => {
210                wait.push(wait_fd(fd as usize, s, type_));
211            }
212            Err(e) => {
213                r_events[i] = handle_event_err(type_, e);
214                i += 1;
215            }
216        }
217    }
218
219    if i == 0 {
220        let ddl = clock.timeout.unwrap();
221        let now = std::time::SystemTime::now();
222        let timeout = ddl.duration_since(now).unwrap_or(Duration::from_secs(0));
223        let sleep = tokio::time::timeout(timeout, wait.select_next_some()).await;
224        if sleep.is_err() {
225            let r_event = &mut r_events[0];
226            r_event.userdata = clock.userdata;
227            r_event.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_CLOCK;
228            mem.write_data(revents_num_ptr, 1)?;
229            return Ok(());
230        }
231
232        let mut connected_fds = vec![];
233
234        let (first, connected_fd) = sleep.unwrap()?;
235        connected_fds.push(connected_fd);
236        r_events[i] = first;
237        i += 1;
238
239        'wait: loop {
240            if i >= nsubscriptions {
241                break 'wait;
242            }
243            futures::select! {
244                v = wait.next() => {
245                    if let Some(v) = v {
246                        let (v,connected_fd) = v?;
247                        connected_fds.push(connected_fd);
248                        r_events[i] = v;
249                        i += 1;
250                    } else {
251                        break 'wait;
252                    }
253                }
254                default => {
255                    break 'wait;
256                }
257            };
258        }
259
260        drop(wait);
261
262        for fd in connected_fds.into_iter().flatten() {
263            if let Ok(socket) = ctx.vfs.get_mut_socket(fd) {
264                socket.state.so_conn_state = ConnectState::Connected;
265            }
266        }
267    }
268
269    mem.write_data(revents_num_ptr, i as u32)?;
270    Ok(())
271}
272
273pub async fn poll_oneoff<M: Memory>(
274    ctx: &mut WasiCtx,
275    mem: &mut M,
276    in_ptr: WasmPtr<__wasi_subscription_t>,
277    out_ptr: WasmPtr<__wasi_event_t>,
278    nsubscriptions: __wasi_size_t,
279    revents_num_ptr: WasmPtr<__wasi_size_t>,
280) -> Result<(), Errno> {
281    log::trace!("poll_oneoff");
282    poll_oneoff_impl(ctx, mem, in_ptr, out_ptr, nsubscriptions, revents_num_ptr).await
283}
284
285async fn poll_oneoff_impl<M: Memory>(
286    ctx: &mut WasiCtx,
287    mem: &mut M,
288    in_ptr: WasmPtr<__wasi_subscription_t>,
289    out_ptr: WasmPtr<__wasi_event_t>,
290    nsubscriptions: __wasi_size_t,
291    revents_num_ptr: WasmPtr<__wasi_size_t>,
292) -> Result<(), Errno> {
293    if nsubscriptions == 0 {
294        return Ok(());
295    }
296
297    let nsubscriptions = nsubscriptions as usize;
298
299    let subs = mem.get_slice(in_ptr, nsubscriptions)?;
300    let prepoll = PrePoll::from_wasi_subscription(subs)?;
301
302    log::trace!("poll_oneoff subs prepoll={:#?}", prepoll);
303
304    match prepoll {
305        PrePoll::OnlyFd(fd_vec) => {
306            poll_only_fd(ctx, mem, out_ptr, nsubscriptions, revents_num_ptr, fd_vec).await?;
307        }
308        PrePoll::ClockAndFd(clock, fd_vec) => {
309            poll_fd_timeout(
310                ctx,
311                mem,
312                out_ptr,
313                nsubscriptions,
314                revents_num_ptr,
315                clock,
316                fd_vec,
317            )
318            .await?;
319        }
320        PrePoll::OnlyClock(clock) => {
321            if let Some(e) = clock.err {
322                let r_event = mem.mut_data(out_ptr)?;
323                r_event.userdata = clock.userdata;
324                r_event.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_CLOCK;
325                r_event.error = e.0;
326                mem.write_data(revents_num_ptr, 1)?;
327                return Ok(());
328            }
329            if let Some(ddl) = clock.timeout {
330                let now = std::time::SystemTime::now();
331                let dur = ddl.duration_since(now).unwrap_or(Duration::from_secs(0));
332                tokio::time::sleep(dur).await;
333                let r_event = mem.mut_data(out_ptr)?;
334                r_event.userdata = clock.userdata;
335                r_event.type_ = __wasi_eventtype_t::__WASI_EVENTTYPE_CLOCK;
336                mem.write_data(revents_num_ptr, 1)?;
337                return Ok(());
338            }
339        }
340    }
341    Ok(())
342}
async_wasi/snapshots/preview_1/async_poll.rs

async_wasi/snapshots/preview_1/
async_poll.rs
