uringmachine 0.19.1 → 0.21.0

data/ext/um/um.c

@@ -1,25 +1,34 @@
-#include <float.h>
 #include "um.h"
-#include "ruby/thread.h"
+#include <float.h>
+#include <ruby/thread.h>
+#include <assert.h>
+#include <poll.h>
+
+#define DEFAULT_ENTRIES 4096
+
+inline void prepare_io_uring_params(struct io_uring_params *params, uint sqpoll_timeout_msec) {
+  memset(params, 0, sizeof(struct io_uring_params));
+  params->flags = IORING_SETUP_SUBMIT_ALL;
+  if (sqpoll_timeout_msec) {
+    params->flags |= IORING_SETUP_SQPOLL;
+    params->sq_thread_idle = sqpoll_timeout_msec;
+  }
+  else
+    params->flags |= IORING_SETUP_COOP_TASKRUN;
+}
 
-void um_setup(VALUE self, struct um *machine) {
+void um_setup(VALUE self, struct um *machine, uint entries, uint sqpoll_timeout_msec) {
   memset(machine, 0, sizeof(struct um));
 
   RB_OBJ_WRITE(self, &machine->self, self);
 
-  unsigned prepared_limit = 4096;
-  unsigned flags = IORING_SETUP_SUBMIT_ALL | IORING_SETUP_COOP_TASKRUN;
-
-  while (1) {
-    int ret = io_uring_queue_init(prepared_limit, &machine->ring, flags);
-    if (likely(!ret)) break;
+  machine->entries = (entries > 0) ? entries : DEFAULT_ENTRIES;
+  machine->sqpoll_mode = !!sqpoll_timeout_msec;
 
-    // if ENOMEM is returned, try with half as much entries
-    if (unlikely(ret == -ENOMEM && prepared_limit > 64))
-      prepared_limit = prepared_limit / 2;
-    else
-      rb_syserr_fail(-ret, strerror(-ret));
-  }
+  struct io_uring_params params;
+  prepare_io_uring_params(&params, sqpoll_timeout_msec);
+  int ret = io_uring_queue_init_params(machine->entries, &machine->ring, &params);
+  if (ret) rb_syserr_fail(-ret, strerror(-ret));
   machine->ring_initialized = 1;
 }
 
@@ -39,11 +48,16 @@ inline void um_teardown(struct um *machine) {
 }
 
 inline struct io_uring_sqe *um_get_sqe(struct um *machine, struct um_op *op) {
+  if (DEBUG) fprintf(stderr, "-> %p um_get_sqe: op->kind=%s unsubmitted=%d pending=%d total=%lu\n",
+    &machine->ring, op ? um_op_kind_name(op->kind) : "NULL", machine->unsubmitted_count,
+    machine->pending_count, machine->total_op_count
+  );
+
   struct io_uring_sqe *sqe;
   sqe = io_uring_get_sqe(&machine->ring);
   if (likely(sqe)) goto done;
 
-  rb_raise(rb_eRuntimeError, "Failed to get SQE");
+  um_raise_internal_error("Failed to get SQE");
 
   // TODO: retry getting SQE?
 
@@ -57,22 +71,78 @@ done:
   sqe->user_data = (long long)op;
   sqe->flags = 0;
   machine->unsubmitted_count++;
-  if (op) machine->pending_count++;
+  if (op) {
+    machine->pending_count++;
+    machine->total_op_count++;
+  }
   return sqe;
 }
 
+struct um_submit_ctx {
+  struct um *machine;
+  int result;
+};
+
+// adapted from liburing/src/queue.c
+static inline bool sq_ring_needs_enter(struct um *machine) {
+  if (machine->sqpoll_mode) {
+	  io_uring_smp_mb();
+  	if (unlikely(IO_URING_READ_ONCE(*machine->ring.sq.kflags) & IORING_SQ_NEED_WAKEUP))
+	  	return true;
+  }
+  return true;
+}
+
+void *um_submit_without_gvl(void *ptr) {
+  struct um_submit_ctx *ctx = ptr;
+  ctx->result = io_uring_submit(&ctx->machine->ring);
+  return NULL;
+}
+
+inline uint um_submit(struct um *machine) {
+  if (DEBUG) fprintf(stderr, "-> %p um_submit: unsubmitted=%d pending=%d total=%lu\n",
+    &machine->ring, machine->unsubmitted_count, machine->pending_count, machine->total_op_count
+  );
+  if (!machine->unsubmitted_count) {
+    if (DEBUG) fprintf(stderr, "<- %p um_submit: no unsubmitted SQEs, early return\n",
+      &machine->ring
+    );
+    return 0;
+  }
+ 
+  struct um_submit_ctx ctx = { .machine = machine };
+  if (sq_ring_needs_enter(machine))
+    rb_thread_call_without_gvl(um_submit_without_gvl, (void *)&ctx, RUBY_UBF_IO, 0);
+  else 
+    ctx.result = io_uring_submit(&machine->ring);
+
+  if (DEBUG) fprintf(stderr, "<- %p um_submit: result=%d\n",
+    &machine->ring, ctx.result
+  );
+
+  if (ctx.result < 0)
+    rb_syserr_fail(-ctx.result, strerror(-ctx.result));
+
+  machine->unsubmitted_count = 0;
+  return ctx.result;
+}
+
 static inline void um_process_cqe(struct um *machine, struct io_uring_cqe *cqe) {
   struct um_op *op = (struct um_op *)cqe->user_data;
+  if (DEBUG) {
+    if (op) fprintf(stderr, "<- %p um_process_cqe: op %p kind %s flags %d cqe_res %d cqe_flags %d pending %d\n",
+      &machine->ring, op, um_op_kind_name(op->kind), op->flags, cqe->res, cqe->flags, machine->pending_count
+    );
+    else fprintf(stderr, "<- %p um_process_cqe: op NULL cqe_res %d cqe_flags %d pending %d\n",
+      &machine->ring, cqe->res, cqe->flags, machine->pending_count
+    );
+  }
   if (unlikely(!op)) return;
 
+
   if (!(cqe->flags & IORING_CQE_F_MORE))
     machine->pending_count--;
 
-  // printf(
-  //   ":process_cqe op %p kind %d flags %d cqe_res %d cqe_flags %d pending %d\n",
-  //   op, op->kind, op->flags, cqe->res, cqe->flags, machine->pending_count
-  // );
-
   if (op->flags & OP_F_FREE_ON_COMPLETE) {
     if (op->flags & OP_F_TRANSIENT)
       um_op_transient_remove(machine, op);
@@ -81,9 +151,10 @@ static inline void um_process_cqe(struct um *machine, struct io_uring_cqe *cqe)
     return;
   }
 
+  op->flags |= OP_F_COMPLETED;
   if (unlikely((cqe->res == -ECANCELED) && (op->flags & OP_F_IGNORE_CANCELED))) return;
+  if (unlikely(op->flags & OP_F_CANCELED)) return;
 
-  op->flags |= OP_F_COMPLETED;
   if (op->flags & OP_F_TRANSIENT)
     um_op_transient_remove(machine, op);
 
@@ -108,6 +179,10 @@ static inline int cq_ring_needs_flush(struct io_uring *ring) {
 }
 
 static inline int um_process_ready_cqes(struct um *machine) {
+  if (DEBUG) fprintf(stderr, "-> %p um_process_ready_cqes: unsubmitted=%d pending=%d total=%lu\n",
+    &machine->ring, machine->unsubmitted_count, machine->pending_count, machine->total_op_count
+  );
+
   unsigned total_count = 0;
 iterate:
   bool overflow_checked = false;
@@ -124,12 +199,21 @@ iterate:
   if (overflow_checked) goto done;
 
   if (cq_ring_needs_flush(&machine->ring)) {
-    io_uring_enter(machine->ring.ring_fd, 0, 0, IORING_ENTER_GETEVENTS, NULL);
+    if (DEBUG) fprintf(stderr, "-> %p io_uring_enter\n", &machine->ring);
+    int ret = io_uring_enter(machine->ring.ring_fd, 0, 0, IORING_ENTER_GETEVENTS, NULL);
+    if (DEBUG) fprintf(stderr, "<- %p io_uring_enter: result=%d\n", &machine->ring, ret);
+    if (ret < 0)
+      rb_syserr_fail(-ret, strerror(-ret));
+
     overflow_checked = true;
     goto iterate;
   }
 
 done:
+  if (DEBUG) fprintf(stderr, "<- %p um_process_ready_cqes: total_processed=%u\n",
+    &machine->ring, total_count
+  );
+
   return total_count;
 }
 
@@ -143,18 +227,33 @@ struct wait_for_cqe_ctx {
 void *um_wait_for_cqe_without_gvl(void *ptr) {
   struct wait_for_cqe_ctx *ctx = ptr;
   if (ctx->machine->unsubmitted_count) {
-    ctx->machine->unsubmitted_count = 0;
+    if (DEBUG) fprintf(stderr, "-> %p io_uring_submit_and_wait_timeout: unsubmitted=%d pending=%d total=%lu\n",
+      &ctx->machine->ring, ctx->machine->unsubmitted_count, ctx->machine->pending_count,
+      ctx->machine->total_op_count
+    );
 
     // Attn: The io_uring_submit_and_wait_timeout will not return -EINTR if
     // interrupted with a signal. We can detect this by testing ctx->cqe for
     // NULL.
     //
     // https://github.com/axboe/liburing/issues/1280
-    int res = io_uring_submit_and_wait_timeout(&ctx->machine->ring, &ctx->cqe, ctx->wait_nr, NULL, NULL);
-    ctx->result = (res > 0 && !ctx->cqe) ? -EINTR : res;
+    int ret = io_uring_submit_and_wait_timeout(&ctx->machine->ring, &ctx->cqe, ctx->wait_nr, NULL, NULL);
+    ctx->machine->unsubmitted_count = 0;
+    if (DEBUG) fprintf(stderr, "<- %p io_uring_submit_and_wait_timeout: result=%d\n",
+      &ctx->machine->ring, ret
+    );
+    ctx->result = (ret > 0 && !ctx->cqe) ? -EINTR : ret;
   }
-  else
+  else {
+    if (DEBUG) fprintf(stderr, "-> %p io_uring_wait_cqes: unsubmitted=%d pending=%d total=%lu\n",
+      &ctx->machine->ring, ctx->machine->unsubmitted_count, ctx->machine->pending_count,
+      ctx->machine->total_op_count
+    );
     ctx->result = io_uring_wait_cqes(&ctx->machine->ring, &ctx->cqe, ctx->wait_nr, NULL, NULL);
+    if (DEBUG) fprintf(stderr, "<- %p io_uring_wait_cqes: result=%d\n",
+      &ctx->machine->ring, ctx->result
+    );
+  }
   return NULL;
 }
 
@@ -201,9 +300,14 @@ inline VALUE process_runqueue_op(struct um *machine, struct um_op *op) {
 }
 
 inline VALUE um_fiber_switch(struct um *machine) {
+  if (DEBUG) fprintf(stderr, "-> %p um_fiber_switch: unsubmitted=%d pending=%d total=%lu\n",
+    &machine->ring, machine->unsubmitted_count, machine->pending_count, machine->total_op_count
+  );
   while (true) {
     struct um_op *op = um_runqueue_shift(machine);
     if (op) {
+      if (unlikely(op->flags & OP_F_RUNQUEUE_SKIP)) continue;
+
       // in case of a snooze, we need to prevent a situation where completions
       // are not processed because the runqueue is never empty. Theoretically,
       // we can still have a situation where multiple fibers are all doing a
@@ -231,16 +335,15 @@ inline VALUE um_fiber_switch(struct um *machine) {
   }
 }
 
-void um_submit_cancel_op(struct um *machine, struct um_op *op) {
+void um_cancel_op(struct um *machine, struct um_op *op) {
   struct io_uring_sqe *sqe = um_get_sqe(machine, NULL);
   io_uring_prep_cancel64(sqe, (long long)op, 0);
 }
 
 inline void um_cancel_and_wait(struct um *machine, struct um_op *op) {
-  um_submit_cancel_op(machine, op);
-  while (true) {
+  um_cancel_op(machine, op);
+  while (!um_op_completed_p(op)) {
     um_fiber_switch(machine);
-    if (um_op_completed_p(op)) break;
   }
 }
 
@@ -261,7 +364,14 @@ inline VALUE um_await(struct um *machine) {
   return ret;
 }
 
-inline void um_prep_op(struct um *machine, struct um_op *op, enum op_kind kind, unsigned flags) {
+VALUE um_wakeup(struct um *machine) {
+  struct io_uring_sqe *sqe = um_get_sqe(machine, NULL);
+  io_uring_prep_nop(sqe);
+  io_uring_submit(&machine->ring);
+  return Qnil;
+}
+
+inline void um_prep_op(struct um *machine, struct um_op *op, enum um_op_kind kind, unsigned flags) {
   memset(op, 0, sizeof(struct um_op));
   op->kind = kind;
   op->flags = flags;
@@ -299,7 +409,7 @@ VALUE um_timeout_complete(VALUE arg) {
   struct op_ctx *ctx = (struct op_ctx *)arg;
 
   if (!um_op_completed_p(ctx->op)) {
-    um_submit_cancel_op(ctx->machine, ctx->op);
+    um_cancel_op(ctx->machine, ctx->op);
     ctx->op->flags |= OP_F_TRANSIENT | OP_F_IGNORE_CANCELED;
     um_op_transient_add(ctx->machine, ctx->op);
   }
@@ -353,12 +463,12 @@ VALUE um_sleep(struct um *machine, double duration) {
   return ret;
 }
 
-inline VALUE um_read(struct um *machine, int fd, VALUE buffer, int maxlen, int buffer_offset) {
+VALUE um_read(struct um *machine, int fd, VALUE buffer, size_t maxlen, ssize_t buffer_offset, __u64 file_offset) {
   struct um_op op;
   um_prep_op(machine, &op, OP_READ, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
   void *ptr = um_prepare_read_buffer(buffer, maxlen, buffer_offset);
-  io_uring_prep_read(sqe, fd, ptr, maxlen, -1);
+  io_uring_prep_read(sqe, fd, ptr, maxlen, file_offset);
 
   VALUE ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op)) {
@@ -366,14 +476,13 @@ inline VALUE um_read(struct um *machine, int fd, VALUE buffer, int maxlen, int b
     ret = INT2NUM(op.result.res);
 
   }
-  RB_GC_GUARD(buffer);
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
   return ret;
 }
 
-inline size_t um_read_raw(struct um *machine, int fd, char *buffer, int maxlen) {
+size_t um_read_raw(struct um *machine, int fd, char *buffer, size_t maxlen) {
   struct um_op op;
   um_prep_op(machine, &op, OP_READ, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
@@ -390,38 +499,47 @@ inline size_t um_read_raw(struct um *machine, int fd, char *buffer, int maxlen)
   return 0;
 }
 
-VALUE um_write(struct um *machine, int fd, VALUE str, int len) {
+VALUE um_write(struct um *machine, int fd, VALUE buffer, size_t len, __u64 file_offset) {
+  const void *base;
+  size_t size;
+  um_get_buffer_bytes_for_writing(buffer, &base, &size);
+  if ((len == (size_t)-1) || (len > size)) len = size;
+  if (unlikely(!len)) return INT2NUM(0);
+
   struct um_op op;
   um_prep_op(machine, &op, OP_WRITE, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
-  const int str_len = RSTRING_LEN(str);
-  if (len > str_len) len = str_len;
 
-  io_uring_prep_write(sqe, fd, RSTRING_PTR(str), len, -1);
+  io_uring_prep_write(sqe, fd, base, len, file_offset);
 
   VALUE ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op))
     ret = INT2NUM(op.result.res);
 
-  RB_GC_GUARD(str);
-
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
   return ret;
 }
 
-VALUE um_write_async(struct um *machine, int fd, VALUE str) {
+VALUE um_write_async(struct um *machine, int fd, VALUE buffer, size_t len, __u64 file_offset) {
+  const void *base;
+  size_t size;
+  um_get_buffer_bytes_for_writing(buffer, &base, &size);
+  if ((len == (size_t)-1) || (len > size)) len = size;
+  if (unlikely(!len)) return INT2NUM(0);
+
   struct um_op *op = um_op_alloc(machine);
   um_prep_op(machine, op, OP_WRITE_ASYNC, OP_F_TRANSIENT | OP_F_FREE_ON_COMPLETE);
   RB_OBJ_WRITE(machine->self, &op->fiber, Qnil);
-  RB_OBJ_WRITE(machine->self, &op->value, str);
+  RB_OBJ_WRITE(machine->self, &op->value, buffer);
   RB_OBJ_WRITE(machine->self, &op->async_op, Qnil);
 
+
   struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  io_uring_prep_write(sqe, fd, RSTRING_PTR(str), RSTRING_LEN(str), -1);
+  io_uring_prep_write(sqe, fd, base, len, file_offset);
   um_op_transient_add(machine, op);
 
-  return str;
+  return buffer;
 }
 
 VALUE um_close(struct um *machine, int fd) {
@@ -497,18 +615,22 @@ VALUE um_connect(struct um *machine, int fd, const struct sockaddr *addr, sockle
   return ret;
 }
 
-VALUE um_send(struct um *machine, int fd, VALUE buffer, int len, int flags) {
+VALUE um_send(struct um *machine, int fd, VALUE buffer, size_t len, int flags) {
   struct um_op op;
   um_prep_op(machine, &op, OP_SEND, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
-  io_uring_prep_send(sqe, fd, RSTRING_PTR(buffer), len, flags);
+
+  const void *base;
+  size_t size;
+  um_get_buffer_bytes_for_writing(buffer, &base, &size);
+  if ((len == (size_t)-1) || (len > size)) len = size;
+
+  io_uring_prep_send(sqe, fd, base, len, flags);
 
   VALUE ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op))
     ret = INT2NUM(op.result.res);
 
-  RB_GC_GUARD(buffer);
-
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
   return ret;
@@ -534,11 +656,12 @@ VALUE um_send_bundle(struct um *machine, int fd, int bgid, VALUE strings) {
   return ret;
 }
 
-VALUE um_recv(struct um *machine, int fd, VALUE buffer, int maxlen, int flags) {
+VALUE um_recv(struct um *machine, int fd, VALUE buffer, size_t maxlen, int flags) {
   struct um_op op;
   um_prep_op(machine, &op, OP_RECV, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
   void *ptr = um_prepare_read_buffer(buffer, maxlen, 0);
+
   io_uring_prep_recv(sqe, fd, ptr, maxlen, flags);
 
   VALUE ret = um_fiber_switch(machine);
@@ -547,8 +670,6 @@ VALUE um_recv(struct um *machine, int fd, VALUE buffer, int maxlen, int flags) {
     ret = INT2NUM(op.result.res);
   }
 
-  RB_GC_GUARD(buffer);
-
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
   return ret;
@@ -588,7 +709,6 @@ VALUE um_getsockopt(struct um *machine, int fd, int level, int opt) {
   VALUE ret = Qnil;
   int value;
 
-#ifdef HAVE_IO_URING_PREP_CMD_SOCK
   struct um_op op;
   um_prep_op(machine, &op, OP_GETSOCKOPT, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
@@ -597,13 +717,6 @@ VALUE um_getsockopt(struct um *machine, int fd, int level, int opt) {
   ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op))
     ret = INT2NUM(value);
-#else
-  socklen_t nvalue = sizeof(value);
-  int res = getsockopt(fd, level, opt, &value, &nvalue);
-  if (res)
-    rb_syserr_fail(errno, strerror(errno));
-  ret = INT2NUM(value);
-#endif
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
@@ -613,7 +726,6 @@ VALUE um_getsockopt(struct um *machine, int fd, int level, int opt) {
 VALUE um_setsockopt(struct um *machine, int fd, int level, int opt, int value) {
   VALUE ret = Qnil;
 
-#ifdef HAVE_IO_URING_PREP_CMD_SOCK
   struct um_op op;
   um_prep_op(machine, &op, OP_SETSOCKOPT, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
@@ -622,12 +734,6 @@ VALUE um_setsockopt(struct um *machine, int fd, int level, int opt, int value) {
   ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op))
     ret = INT2NUM(op.result.res);
-#else
-  int res = setsockopt(fd, level, opt, &value, sizeof(value));
-  if (res)
-    rb_syserr_fail(errno, strerror(errno));
-  ret = INT2NUM(0);
-#endif
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
@@ -694,13 +800,121 @@ VALUE um_poll(struct um *machine, int fd, unsigned mask) {
   return ret;
 }
 
-VALUE um_waitpid(struct um *machine, int pid, int options) {
+static inline void prepare_select_poll_ops(struct um *machine, uint *idx, struct um_op *ops, VALUE fds, uint len, uint flags, uint event) {
+  for (uint i = 0; i < len; i++) {
+    struct um_op *op = ops + ((*idx)++);
+    um_prep_op(machine, op, OP_POLL, flags | OP_F_IGNORE_CANCELED);
+    struct io_uring_sqe *sqe = um_get_sqe(machine, op);
+    VALUE fd = rb_ary_entry(fds, i);
+    RB_OBJ_WRITE(machine->self, &op->value, fd);
+    io_uring_prep_poll_add(sqe, NUM2INT(fd), event);
+  }
+}
+
+VALUE um_select_single(struct um *machine, VALUE rfds, VALUE wfds, VALUE efds, uint rfds_len, uint wfds_len, uint efds_len) {
   struct um_op op;
-  um_prep_op(machine, &op, OP_WAITPID, 0);
+  uint idx = 0;
+  if (rfds_len)
+    prepare_select_poll_ops(machine, &idx, &op, rfds, rfds_len, OP_F_SELECT_POLLIN, POLLIN);
+  else if (wfds_len)
+    prepare_select_poll_ops(machine, &idx, &op, wfds, wfds_len, OP_F_SELECT_POLLOUT, POLLOUT);
+  else if (efds_len)
+    prepare_select_poll_ops(machine, &idx, &op, efds, efds_len, OP_F_SELECT_POLLPRI, POLLPRI);
+  assert(idx == 1);
+
+  VALUE ret = um_fiber_switch(machine);
+  um_check_completion(machine, &op);
+  RAISE_IF_EXCEPTION(ret);
+
+  if (op.flags & OP_F_SELECT_POLLIN)
+    return rb_ary_new3(3, rb_ary_new3(1, ret), rb_ary_new(), rb_ary_new());
+  else if (op.flags & OP_F_SELECT_POLLOUT)
+    return rb_ary_new3(3, rb_ary_new(), rb_ary_new3(1, ret), rb_ary_new());
+  else
+    return rb_ary_new3(3, rb_ary_new(), rb_ary_new(), rb_ary_new3(1, ret));
+
+  RB_GC_GUARD(ret);
+}
+
+VALUE um_select(struct um *machine, VALUE rfds, VALUE wfds, VALUE efds) {
+  uint rfds_len = RARRAY_LEN(rfds);
+  uint wfds_len = RARRAY_LEN(wfds);
+  uint efds_len = RARRAY_LEN(efds);
+  uint total_len = rfds_len + wfds_len + efds_len;
+  if (total_len == 1)
+    return um_select_single(machine, rfds, wfds, efds, rfds_len, wfds_len, efds_len);
+
+  if (unlikely(!total_len))
+    return rb_ary_new3(3, rb_ary_new(), rb_ary_new(), rb_ary_new());
+
+  struct um_op *ops = malloc(sizeof(struct um_op) * total_len);
+  uint idx = 0;
+  prepare_select_poll_ops(machine, &idx, ops, rfds, rfds_len, OP_F_SELECT_POLLIN, POLLIN);
+  prepare_select_poll_ops(machine, &idx, ops, wfds, wfds_len, OP_F_SELECT_POLLOUT, POLLOUT);
+  prepare_select_poll_ops(machine, &idx, ops, efds, efds_len, OP_F_SELECT_POLLPRI, POLLPRI);
+  assert(idx == total_len);
+
+  VALUE ret = um_fiber_switch(machine);
+  if (unlikely(um_value_is_exception_p(ret))) {
+    free(ops);
+    um_raise_exception(ret);
+  }
+
+  VALUE rfds_out = rb_ary_new();
+  VALUE wfds_out = rb_ary_new();
+  VALUE efds_out = rb_ary_new();
+
+  int error_code = 0;
+  uint pending = total_len;
+  for (uint i = 0; i < total_len; i++) {
+    if (um_op_completed_p(&ops[i])) {
+      ops[i].flags |= OP_F_RUNQUEUE_SKIP;
+      pending--;
+
+      if (unlikely((ops[i].result.res < 0) && !error_code)) {
+        error_code = ops[i].result.res;
+      }
+      else {
+        if (ops[i].flags & OP_F_SELECT_POLLIN)  rb_ary_push(rfds_out, ops[i].value);
+        if (ops[i].flags & OP_F_SELECT_POLLOUT) rb_ary_push(wfds_out, ops[i].value);
+        if (ops[i].flags & OP_F_SELECT_POLLPRI) rb_ary_push(efds_out, ops[i].value);
+      }
+    }
+    else {
+      ops[i].flags |= OP_F_CANCELED;
+      um_cancel_op(machine, &ops[i]);
+    }
+  }
+
+  while (pending) {
+    um_wait_for_and_process_ready_cqes(machine, 0);
+
+    for (uint i = 0; i < total_len; i++) {
+      struct um_op *op = ops + i;
+      if (op->flags & OP_F_CANCELED && um_op_completed_p(op)) {
+        pending--;
+      }
+    }
+  }
+  free(ops);
+
+  if (error_code)
+    um_raise_on_error_result(error_code);
+
+  return rb_ary_new3(3, rfds_out, wfds_out, efds_out);
+
+  RB_GC_GUARD(rfds_out);
+  RB_GC_GUARD(wfds_out);
+  RB_GC_GUARD(efds_out);
+}
+
+VALUE um_waitid(struct um *machine, int idtype, int id, int options) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_WAITID, 0);
   struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
 
   siginfo_t infop;
-  io_uring_prep_waitid(sqe, pid == 0 ? P_ALL : P_PID, pid, &infop, options, 0);
+  io_uring_prep_waitid(sqe, idtype, id, &infop, options, 0);
 
   VALUE ret = um_fiber_switch(machine);
   if (um_check_completion(machine, &op))
@@ -709,9 +923,31 @@ VALUE um_waitpid(struct um *machine, int pid, int options) {
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
 
-  return rb_ary_new_from_args(2, INT2NUM(infop.si_pid), INT2NUM(infop.si_status));
+  return rb_ary_new_from_args(
+    3, INT2NUM(infop.si_pid), INT2NUM(infop.si_status), INT2NUM(infop.si_code)
+  );
 }
 
+#ifdef HAVE_RB_PROCESS_STATUS_NEW
+VALUE um_waitid_status(struct um *machine, int idtype, int id, int options) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_WAITID, 0);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+
+  siginfo_t infop;
+  io_uring_prep_waitid(sqe, idtype, id, &infop, options | WNOWAIT, 0);
+
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
+
+  RAISE_IF_EXCEPTION(ret);
+  RB_GC_GUARD(ret);
+
+  return rb_process_status_new(infop.si_pid, (infop.si_status & 0xff) << 8, 0);
+}
+#endif
+
 #define hash_set(h, sym, v) rb_hash_aset(h, ID2SYM(rb_intern(sym)), v)
 
 VALUE statx_to_hash(struct statx *stat) {