wasip3: Implement nonblocking UDP I/O (#775)

This commit implements nonblocking I/O for UDP sockets, specifically
with `send` and `recv`. This additionally fills out enough of `poll` to
get various tests doing nonblocking I/O to work. The main intention of
this commit is to do the first foray into implementing nonblocking I/O
in wasi-libc for wasip3. Broadly-speaking nonblocking I/O is in two
categories: one being UDP and the other being `stream<u8>`-based
reads/writes. This implements only the UDP half, and defers
`stream<u8>` (and thus TCP for example) to a future commit.

The implementation is relatively gnarly given the wasip3 APIs we
currently have, but this was expected. I've done my best to be as
careful as I can in terms of managing state machines but I'm pretty
likely to inevitably get something wrong, so this should be expected to
be something that's iterated on.

Author: alexcrichton

expected/wasm32-wasip3/defined-symbols.txt

@@ -329,6 +329,9 @@ __wasilibc_nocwd_symlinkat
 __wasilibc_nocwd_utimensat
 __wasilibc_open_nomode
 __wasilibc_parse_port
+__wasilibc_poll_add
+__wasilibc_poll_ready
+__wasilibc_poll_waitable
 __wasilibc_populate_preopens
 __wasilibc_pthread_self
 __wasilibc_random

libc-bottom-half/cloudlibc/src/libc/poll/poll.c

@@ -16,26 +16,6 @@ static int poll_impl(struct pollfd *fds, size_t nfds, int timeout) {
   __wasi_subscription_t subscriptions[maxevents];
   size_t nsubscriptions = 0;
 
-  // POLLPRI (exceptional/out-of-band data) is not supported in WASI.
-  // If all requested events across all fds are POLLPRI, return ENOSYS.
-  // Otherwise, remove POLLPRI so it never fires (same as no OOB data).
-  {
-    bool has_pri_only = false;
-    bool has_non_pri = false;
-    for (size_t i = 0; i < nfds; ++i) {
-      if (fds[i].fd < 0 || fds[i].events == 0)
-        continue;
-      if (fds[i].events & ~POLLPRI)
-        has_non_pri = true;
-      else
-        has_pri_only = true;
-    }
-    if (has_pri_only && !has_non_pri) {
-      errno = ENOSYS;
-      return -1;
-    }
-  }
-
   for (size_t i = 0; i < nfds; ++i) {
     struct pollfd *pollfd = &fds[i];
     if (pollfd->fd < 0)
@@ -105,12 +85,6 @@ static int poll_impl(struct pollfd *fds, size_t nfds, int timeout) {
     return -1;
   }
 
-  // Clear revents fields.
-  for (size_t i = 0; i < nfds; ++i) {
-    struct pollfd *pollfd = &fds[i];
-    pollfd->revents = 0;
-  }
-
   // Set revents fields.
   for (size_t i = 0; i < nevents; ++i) {
     const __wasi_event_t *event = &events[i];
@@ -196,29 +170,6 @@ void __wasilibc_poll_ready(poll_state_t *state, short events) {
 }
 
 static int poll_impl(struct pollfd *fds, size_t nfds, int timeout) {
-  for (size_t i = 0; i < nfds; ++i) {
-    fds[i].revents = 0;
-  }
-
-  // POLLPRI (exceptional/out-of-band data) is not supported in WASI.
-  // If all requested events across all fds are exclusively POLLPRI,
-  // return ENOSYS. Otherwise, strip POLLPRI and proceed.
-  {
-    bool has_pri_only = false;
-    bool has_non_pri = false;
-    for (size_t i = 0; i < nfds; ++i) {
-      if (fds[i].fd < 0 || fds[i].events == 0)
-        continue;
-      if (fds[i].events & ~POLLPRI)
-        has_non_pri = true;
-      else
-        has_pri_only = true;
-    }
-    if (has_pri_only && !has_non_pri) {
-      errno = ENOSYS;
-      return -1;
-    }
-  }
 
   size_t max_pollables = (2 * nfds) + 1;
   state_t states[max_pollables];
@@ -330,18 +281,237 @@ static int poll_impl(struct pollfd *fds, size_t nfds, int timeout) {
 
 #elif defined(__wasip3__)
 
+typedef struct {
+  struct pollfd *pollfd;
+  uint32_t waitable;
+  // Callback/data pointer to invoke when `waitable` receives an event, and that
+  // runs any completion logic and figures out what to do next.
+  poll_ready_t ready;
+  void *ready_data;
+} state_t;
+
+struct poll_state_t {
+  state_t *states;
+  size_t len;
+  size_t cap;
+
+  wasip3_waitable_set_t set;
+
+  int event_count;
+  struct pollfd *pollfd;
+};
+
+int __wasilibc_poll_add(poll_state_t *state, uint32_t waitable,
+                        poll_ready_t ready, void *ready_data) {
+  if (state->len >= state->cap) {
+    errno = ENOMEM;
+    return -1;
+  }
+  state->states[state->len].pollfd = state->pollfd;
+  state->states[state->len].waitable = waitable;
+  state->states[state->len].ready = ready;
+  state->states[state->len].ready_data = ready_data;
+  wasip3_waitable_join(waitable, state->set);
+  state->len += 1;
+  return 0;
+}
+
+void __wasilibc_poll_ready(poll_state_t *state, short events) {
+  if (events != 0) {
+    if (state->pollfd->revents == 0) {
+      ++state->event_count;
+    }
+    state->pollfd->revents |= events;
+  }
+}
+
 static int poll_impl(struct pollfd *fds, size_t nfds, int timeout) {
-  (void) fds;
-  (void) nfds;
-  (void) timeout;
-  errno = ENOTSUP;
-  return -1;
+  // TODO(wasip3) this calculation of 2n+1 is a coarse approximation but not
+  // necessarily accurate. This worked well for wasip2 but this should be
+  // revisited for wasip3. Maybe make this a dynamically allocated array?
+  size_t max_pollables = (2 * nfds) + 1;
+  state_t states[max_pollables];
+
+  poll_state_t state;
+  state.set = wasip3_waitable_set_new();
+  state.states = states;
+  state.len = 0;
+  state.cap = max_pollables;
+  state.event_count = 0;
+  state.pollfd = fds;
+
+  int ret = -1;
+  wasip3_subtask_t timeout_subtask = 0;
+
+  for (size_t i = 0; i < nfds; ++i) {
+    struct pollfd *pollfd = fds + i;
+    if (pollfd->fd < 0)
+      continue;
+    state.pollfd = pollfd;
+    descriptor_table_entry_t *entry = descriptor_table_get_ref(pollfd->fd);
+    if (!entry) {
+      errno = EBADF;
+      goto out;
+    }
+
+    // If this descriptor has a custom registration function then
+    // use that exclusively.
+    if (entry->vtable->poll_register) {
+      if (entry->vtable->poll_register(entry->data, &state, pollfd->events) < 0)
+        goto out;
+      continue;
+    }
+
+    // Strip POLLPRI, it is never reported in WASI (no OOB data).
+    short events = pollfd->events & ~POLLPRI;
+
+    // Without a custom registration handle read/write readiness
+    // below, but everything else is unsupported.
+    if (events & ~(POLLRDNORM | POLLWRNORM)) {
+      errno = EOPNOTSUPP;
+      goto out;
+    }
+
+
+    if (events & POLLRDNORM) {
+      // TODO(wasip3): use `get_read_stream` and use that I/O to do something.
+      // Requires a lot more integration with nonblocking I/O to get that
+      // working.
+      errno = EOPNOTSUPP;
+      goto out;
+    }
+
+    if (events & POLLWRNORM) {
+      // TODO(wasip3): use `get_write_stream` to implement this (see
+      // `POLLRDNORM` above).
+      errno = EOPNOTSUPP;
+      goto out;
+    }
+  }
+
+  // If something ended up being ready during registration then zero out the
+  // timeout and continue proceeding as if this is a non-blocking call.
+  if (state.event_count > 0)
+    timeout = 0;
+
+  // If a timeout is specified then spawn a subtask using
+  // `monotonic-clock#wait-for`. This subtask is then added to the
+  // waitable-set.  If the timeout completes immediately then pretend like
+  // there's no timeout and keep going.
+  //
+  // Note that in `out`, the exit of this function, the subtask is cleaned up
+  // if it's still in-progress.
+  if (timeout > 0) {
+    wasip3_subtask_status_t status = monotonic_clock_wait_for(timeout * 1000000);
+    if (WASIP3_SUBTASK_STATE(status) == WASIP3_SUBTASK_RETURNED) {
+      timeout = 0;
+    } else {
+      timeout_subtask = WASIP3_SUBTASK_HANDLE(status);
+      wasip3_waitable_join(timeout_subtask, state.set);
+    }
+  }
+
+  // Perform the actual blocking operation, if applicable. If the timeout is
+  // still nonzero at this point then it means that it's time to actually
+  // block. That handles the case where the timeout is negative as no subtask
+  // was spawned and this will wait indefinitely. It also handles the case of a
+  // timeout where a subtask was spawned and it's in the set being waited on.
+  //
+  // If the timeout is 0, however, then just perform a poll-style operation to
+  // see what happened.
+  wasip3_event_t event;
+  if (timeout != 0) {
+    wasip3_waitable_set_wait(state.set, &event);
+  } else {
+    wasip3_waitable_set_poll(state.set, &event);
+  }
+  while (event.event != WASIP3_EVENT_NONE) {
+    // If this event is for the timeout subtask then that's handled directly here
+    // as it's not related to any fds. Upon receiving the terminal status of
+    // the subtask it's no longer candidate for cancellation so it's immediately
+    // dropped and zero'd out to avoid processing in the `out` label below.
+    if (event.waitable == timeout_subtask) {
+      assert(event.event == WASIP3_EVENT_SUBTASK);
+      assert(event.code == WASIP3_SUBTASK_RETURNED);
+      wasip3_subtask_drop(timeout_subtask);
+      timeout_subtask = 0;
+    } else {
+      // If this event is not for the timeout subtask then it's something that
+      // was added to the waitable-set. That's processed here directly. Search
+      // through the list of `state_t` waitables, find the right one, and invoke
+      // its `ready` callback. The `ready` callback will dictate what to do
+      // with `event` and will call `__wasilibc_poll_ready` as appropriate.
+      //
+      // TODO: having to search through the list isn't great, but short of
+      // having some sort of hash table or map going from waitable->state_t
+      // that's the best that can be done right. At least `poll` as a POSIX
+      // interface is expected to take linear time, right?
+      for (size_t i = 0; i < state.len; i++) {
+        state_t *p = &state.states[i];
+        if (p->waitable != event.waitable)
+          continue;
+        state.pollfd = p->pollfd;
+        // Clear the `waitable` since this event has fired and it doesn't need
+        // to be join'd to 0 below. Then invoke the custom callback originally
+        // added for this which will handle any necessary completion logic and
+        // updating `state.pollfd` with various events.
+        p->waitable = 0;
+        p->ready(p->ready_data, &state, &event);
+      }
+    }
+
+    // After this event has finished being processed try to get another event.
+    // This will drain the waitable-set of all ready events in one invocation of
+    // this function and avoid unncessary setup/teardown in `poll`, in theory.
+    wasip3_waitable_set_poll(state.set, &event);
+  }
+
+  ret = state.event_count;
+
+out:
+  if (timeout_subtask != 0) {
+    wasip3_subtask_cancel(timeout_subtask);
+    wasip3_subtask_drop(timeout_subtask);
+  }
+  for (size_t i = 0; i < state.len; i++) {
+    uint32_t waitable = state.states[i].waitable;
+    if (waitable)
+      wasip3_waitable_join(waitable, 0);
+  }
+  wasip3_waitable_set_drop(state.set);
+
+  return ret;
 }
 
 #else
 # error "Unknown WASI version"
 #endif
 
+// POLLPRI (exceptional/out-of-band data) is not supported in WASI.
+// If all requested events across all fds are exclusively POLLPRI,
+// return ENOSYS. Otherwise, strip POLLPRI and proceed.
+static int validate_something_not_pollpri(struct pollfd *fds, size_t nfds) {
+  bool has_pri_only = false;
+  bool has_non_pri = false;
+  for (size_t i = 0; i < nfds; ++i) {
+    if (fds[i].fd < 0 || fds[i].events == 0)
+      continue;
+    if (fds[i].events & ~POLLPRI)
+      has_non_pri = true;
+    else
+      has_pri_only = true;
+  }
+  if (has_pri_only && !has_non_pri) {
+    errno = ENOSYS;
+    return -1;
+  }
+  return 0;
+}
+
 int poll(struct pollfd* fds, nfds_t nfds, int timeout) {
-    return poll_impl(fds, nfds, timeout);
+  if (validate_something_not_pollpri(fds, nfds) < 0)
+    return -1;
+  for (size_t i = 0; i < nfds; ++i)
+    fds[i].revents = 0;
+  return poll_impl(fds, nfds, timeout);
 }

libc-bottom-half/headers/private/wasi/poll.h

@@ -6,6 +6,7 @@
 #ifndef __wasip1__
 
 #include <poll.h>
+#include <wasi/api.h>
 
 /// Opaque state that the `poll` function manages for itself.
 typedef struct poll_state_t poll_state_t;
@@ -18,8 +19,6 @@ void __wasilibc_poll_ready(poll_state_t *state, short events);
 
 #ifdef __wasip2__
 
-#include <wasi/wasip2.h>
-
 /// Adds the `pollable` to the `state` provided.
 ///
 /// This will add a `pollable` to get passed to `wasi:io/poll.poll`. If that
@@ -56,6 +55,20 @@ __wasilibc_poll_add_output_stream(poll_state_t *state,
 
 #endif // __wasip2__
 
+#ifdef __wasip3__
+
+typedef void (*poll_ready_t)(void *data, poll_state_t *state,
+                             wasip3_event_t *event);
+
+/// Adds the `waitable` to the poll `state`.
+///
+/// When `waitable` receives an event it'll invoke `ready` with `ready_data`,
+/// the poll state, and the event that happened.
+int __wasilibc_poll_add(poll_state_t *state, uint32_t waitable,
+                        poll_ready_t ready, void *ready_data);
+
+#endif // __wasip3__
+
 #endif // not(__wasip1__)
 
 #endif // WASI_POLL_H

libc-bottom-half/headers/private/wasi/wasip3_block.h

@@ -83,6 +83,12 @@ static inline size_t __wasilibc_stream_block_on(wasip3_waitable_status_t status,
   return ret;
 }
 
+/// Polls the `waitable` provide, filling in `event`, if any.
+///
+/// This will call `waitable-set.poll` with `waitable` and the `event` is
+/// whatever that returns.
+void __wasilibc_poll_waitable(uint32_t waitable, wasip3_event_t *event);
+
 #endif // __wasip3__
 
 #endif // WASI_WASIP3_BLOCK_H