uringmachine 0.4 → 0.5.1

data/ext/um/um.c

@@ -1,24 +1,13 @@
 #include "um.h"
 #include "ruby/thread.h"
 
-void um_setup(struct um *machine) {
-  machine->ring_initialized = 0;
-  machine->unsubmitted_count = 0;
-  machine->buffer_ring_count = 0;
-  machine->pending_count = 0;
-  machine->runqueue_head = NULL;
-  machine->runqueue_tail = NULL;
-  machine->op_freelist = NULL;
-  machine->result_freelist = NULL;
+void um_setup(VALUE self, struct um *machine) {
+  memset(machine, 0, sizeof(struct um));
+
+  RB_OBJ_WRITE(self, &machine->self, self);
 
   unsigned prepared_limit = 4096;
-  unsigned flags = 0;
-  #ifdef HAVE_IORING_SETUP_SUBMIT_ALL
-  flags |= IORING_SETUP_SUBMIT_ALL;
-  #endif
-  #ifdef HAVE_IORING_SETUP_COOP_TASKRUN
-  flags |= IORING_SETUP_COOP_TASKRUN;
-  #endif
+  unsigned flags = IORING_SETUP_SUBMIT_ALL | IORING_SETUP_COOP_TASKRUN;
 
   while (1) {
     int ret = io_uring_queue_init(prepared_limit, &machine->ring, flags);
@@ -45,11 +34,10 @@ inline void um_teardown(struct um *machine) {
   io_uring_queue_exit(&machine->ring);
   machine->ring_initialized = 0;
 
-  um_free_op_linked_list(machine, machine->op_freelist);
-  um_free_op_linked_list(machine, machine->runqueue_head);
+  um_free_buffer_linked_list(machine);
 }
 
-static inline struct io_uring_sqe *um_get_sqe(struct um *machine, struct um_op *op) {
+inline struct io_uring_sqe *um_get_sqe(struct um *machine, struct um_op *op) {
   struct io_uring_sqe *sqe;
   sqe = io_uring_get_sqe(&machine->ring);
   if (likely(sqe)) goto done;
@@ -68,92 +56,44 @@ done:
   sqe->user_data = (long long)op;
   sqe->flags = 0;
   machine->unsubmitted_count++;
+  if (op) machine->pending_count++;
   return sqe;
 }
 
-struct wait_for_cqe_ctx {
-  struct um *machine;
-  struct io_uring_cqe *cqe;
-  int result;
-};
-
-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;
-    ctx->result = io_uring_submit_and_wait_timeout(&ctx->machine->ring, &ctx->cqe, 1, NULL, NULL);
-  }
-  else
-    ctx->result = io_uring_wait_cqe(&ctx->machine->ring, &ctx->cqe);
-  return NULL;
-}
-
-inline void um_handle_submitted_op_cqe_single(struct um *machine, struct um_op *op, struct io_uring_cqe *cqe) {
-  op->cqe_result = cqe->res;
-  op->cqe_flags = cqe->flags;
-  op->state = OP_completed;
-  um_runqueue_push(machine, op);
-}
-
-inline void um_handle_submitted_op_cqe_multi(struct um *machine, struct um_op *op, struct io_uring_cqe *cqe) {
-  if (!op->results_head) {
-    // if no results are ready yet, schedule the corresponding fiber
-    struct um_op *op2 = um_op_checkout(machine);
-    op2->state = OP_schedule;
-    op2->fiber = op->fiber;
-    op2->resume_value = Qnil;
-    um_runqueue_push(machine, op2);
-  }
-  um_op_result_push(machine, op, cqe->res, cqe->flags);
-
-  if (!(cqe->flags & IORING_CQE_F_MORE))
-    op->state = OP_completed;
-}
-
-inline void um_process_cqe(struct um *machine, struct io_uring_cqe *cqe) {
+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 (unlikely(!op)) return;
 
-  // if (op->is_multishot)
-  //   printf("process_cqe %p state: %d result: %d flags: %d (%d)\n", op, op->state, cqe->res, cqe->flags, (cqe->flags & IORING_CQE_F_MORE));
+  if (!(cqe->flags & IORING_CQE_F_MORE))
+    machine->pending_count--;
 
-  switch (op->state) {
-    case OP_submitted:
-      if (unlikely(cqe->res == -ECANCELED)) {
-        um_op_checkin(machine, op);
-        break;
-      }
-      if (op->is_multishot)
-        um_handle_submitted_op_cqe_multi(machine, op, cqe);
-      else
-        um_handle_submitted_op_cqe_single(machine, op, cqe);
-      break;
-    case OP_abandonned:
-      // op has been abandonned by the I/O method, so we need to cleanup (check
-      // the op in to the free list).
-      um_op_checkin(machine, op);
-      break;
-    default:
-      // TODO: invalid state, should raise!
-  }
-}
+  // 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
+  // );
 
-static inline void um_wait_for_and_process_cqe(struct um *machine) {
-  struct wait_for_cqe_ctx ctx = {
-    .machine = machine,
-    .cqe = NULL
-  };
+  if (unlikely((cqe->res == -ECANCELED) && (op->flags & OP_F_IGNORE_CANCELED))) return;
 
-  rb_thread_call_without_gvl(um_wait_for_cqe_without_gvl, (void *)&ctx, RUBY_UBF_IO, 0);
-  if (unlikely(ctx.result < 0)) {
-    rb_syserr_fail(-ctx.result, strerror(-ctx.result));
+  op->flags |= OP_F_COMPLETED;
+  if (op->flags & OP_F_TRANSIENT)
+    um_op_transient_remove(machine, op);
+
+  if (op->flags & OP_F_MULTISHOT) {
+    um_op_multishot_results_push(machine, op, cqe->res, cqe->flags);
+    if (op->multishot_result_count > 1)
+      return;
   }
-  io_uring_cqe_seen(&machine->ring, ctx.cqe);
-  um_process_cqe(machine, ctx.cqe);
+  else {
+    op->result.res    = cqe->res;
+    op->result.flags  = cqe->flags;
+  }
+
+  if (!(op->flags & OP_F_ASYNC))
+    um_runqueue_push(machine, op);
 }
 
 // copied from liburing/queue.c
-static inline bool cq_ring_needs_flush(struct io_uring *ring) {
+static inline int cq_ring_needs_flush(struct io_uring *ring) {
   return IO_URING_READ_ONCE(*ring->sq.kflags) & IORING_SQ_CQ_OVERFLOW;
 }
 
@@ -183,101 +123,117 @@ done:
   return total_count;
 }
 
-static inline void um_wait_for_and_process_ready_cqes(struct um *machine) {
-  um_wait_for_and_process_cqe(machine);
-  um_process_ready_cqes(machine);
+struct wait_for_cqe_ctx {
+  struct um *machine;
+  struct io_uring_cqe *cqe;
+  int result;
+};
+
+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;
+
+    // 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, 1, NULL, NULL);
+    ctx->result = (res > 0 && !ctx->cqe) ? -EINTR : res;
+  }
+  else
+    ctx->result = io_uring_wait_cqe(&ctx->machine->ring, &ctx->cqe);
+  return NULL;
 }
 
-inline VALUE um_fiber_switch(struct um *machine) {
-  struct um_op *op = 0;
-  unsigned int first_iteration = 1;
-loop:
-  // in the case where:
-  // - first time through loop
-  // - there are SQEs waiting to be submitted
-  // - the runqueue head references the current fiber
-  // we need to submit events and check completions without blocking
-  if (
-    unlikely(
-      first_iteration && machine->unsubmitted_count &&
-      machine->runqueue_head &&
-      machine->runqueue_head->fiber == rb_fiber_current()
-    )
-  ) {
-    io_uring_submit(&machine->ring);
+static inline void um_wait_for_and_process_ready_cqes(struct um *machine) {
+  struct wait_for_cqe_ctx ctx = { .machine = machine, .cqe = NULL };
+  rb_thread_call_without_gvl(um_wait_for_cqe_without_gvl, (void *)&ctx, RUBY_UBF_IO, 0);
+
+  if (unlikely(ctx.result < 0 && ctx.result != -EINTR))
+    rb_syserr_fail(-ctx.result, strerror(-ctx.result));
+
+  if (ctx.cqe) {
+    um_process_cqe(machine, ctx.cqe);
+    io_uring_cq_advance(&machine->ring, 1);
     um_process_ready_cqes(machine);
   }
-  first_iteration = 0;
+}
 
-  op = um_runqueue_shift(machine);
-  if (op) {
-    VALUE resume_value = op->resume_value;
-    if (op->state == OP_schedule)
-      um_op_checkin(machine, op);
+inline VALUE process_runqueue_op(struct um *machine, struct um_op *op) {
+  VALUE fiber = op->fiber;
+  VALUE value = op->value;
 
-    // the resume value is disregarded, we pass the fiber itself
-    VALUE v = rb_fiber_transfer(op->fiber, 1, &resume_value);
-    return v;
-  }
+  if (unlikely(op->flags & OP_F_TRANSIENT))
+    um_op_free(machine, op);
 
-  um_wait_for_and_process_ready_cqes(machine);
-  goto loop;
+  return rb_fiber_transfer(fiber, 1, &value);
 }
 
-static inline void um_cancel_op(struct um *machine, struct um_op *op) {
+inline VALUE um_fiber_switch(struct um *machine) {
+  while (true) {
+    struct um_op *op = um_runqueue_shift(machine);
+    if (op)
+      return process_runqueue_op(machine, op);
+
+    um_wait_for_and_process_ready_cqes(machine);
+  }
+}
+
+void um_submit_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);
 }
 
-static inline VALUE um_await_op(struct um *machine, struct um_op *op, __s32 *result, __u32 *flags) {
-  op->fiber = rb_fiber_current();
-  VALUE v = um_fiber_switch(machine);
-  int is_exception = um_value_is_exception_p(v);
-
-  if (unlikely(is_exception && op->state == OP_submitted)) {
-    um_cancel_op(machine, op);
-    op->state = OP_abandonned;
+inline void um_cancel_and_wait(struct um *machine, struct um_op *op) {
+  um_submit_cancel_op(machine, op);
+  while (true) {
+    um_fiber_switch(machine);
+    if (um_op_completed_p(op)) break;
   }
-  else {
-    // We copy over the CQE result and flags, since the op is immediately
-    // checked in.
-    if (result) *result = op->cqe_result;
-    if (flags) *flags = op->cqe_flags;
+}
+
+inline int um_check_completion(struct um *machine, struct um_op *op) {
+  if (!um_op_completed_p(op)) {
+    um_cancel_and_wait(machine, op);
+    return 0;
   }
 
-  if (unlikely(is_exception)) um_raise_exception(v);
-  return v;
+  um_raise_on_error_result(op->result.res);
+  return 1;
 }
 
 inline VALUE um_await(struct um *machine) {
   VALUE v = um_fiber_switch(machine);
-  return um_value_is_exception_p(v) ? um_raise_exception(v) : v;
+  return raise_if_exception(v);
 }
 
-inline void um_schedule(struct um *machine, VALUE fiber, VALUE value) {
-  struct um_op *op = um_op_checkout(machine);
-  op->state = OP_schedule;
-  op->fiber = fiber;
-  op->resume_value = value;
-  um_runqueue_push(machine, op);
+inline void um_prep_op(struct um *machine, struct um_op *op, enum op_kind kind) {
+  memset(op, 0, sizeof(struct um_op));
+  op->kind = kind;
+  switch (kind) {
+    case OP_ACCEPT_MULTISHOT:
+    case OP_READ_MULTISHOT:
+    case OP_RECV_MULTISHOT:
+      op->flags |= OP_F_MULTISHOT;
+    default:
+  }
+  RB_OBJ_WRITE(machine->self, &op->fiber, rb_fiber_current());
+  op->value = Qnil;
 }
 
-inline void um_interrupt(struct um *machine, VALUE fiber, VALUE value) {
-  struct um_op *op = um_runqueue_find_by_fiber(machine, fiber);
-  if (op) {
-    op->state = OP_cancelled;
-    op->resume_value = value;
-  }
-  else {
-    op = um_op_checkout(machine);
-    op->state = OP_schedule;
-    op->fiber = fiber;
-    op->resume_value = value;
-    um_runqueue_unshift(machine, op);
-  }
+inline void um_schedule(struct um *machine, VALUE fiber, VALUE value) {
+  struct um_op *op = um_op_alloc(machine);
+  memset(op, 0, sizeof(struct um_op));
+  op->kind = OP_SCHEDULE;
+  op->flags = OP_F_TRANSIENT;
+  RB_OBJ_WRITE(machine->self, &op->fiber, fiber);
+  RB_OBJ_WRITE(machine->self, &op->value, value);
+  um_runqueue_push(machine, op);
 }
 
-struct op_ensure_ctx {
+struct op_ctx {
   struct um *machine;
   struct um_op *op;
   int fd;
@@ -285,20 +241,18 @@ struct op_ensure_ctx {
 
   void *read_buf;
   int read_maxlen;
+  struct __kernel_timespec ts;
+  int flags;
 };
 
 VALUE um_timeout_ensure(VALUE arg) {
-  struct op_ensure_ctx *ctx = (struct op_ensure_ctx *)arg;
+  struct op_ctx *ctx = (struct op_ctx *)arg;
 
-  if (ctx->op->state == OP_submitted) {
-    // A CQE has not yet been received, we cancel the timeout and abandon the op
-    // (it will be checked in upon receiving the -ECANCELED CQE)
-    um_cancel_op(ctx->machine, ctx->op);
-    ctx->op->state == OP_abandonned;
+  if (!um_op_completed_p(ctx->op)) {
+    um_submit_cancel_op(ctx->machine, ctx->op);
+    ctx->op->flags |= OP_F_TRANSIENT | OP_F_IGNORE_CANCELED;
+    um_op_transient_add(ctx->machine, ctx->op);
   }
-  else
-    // completed, so can be checked in
-    um_op_checkin(ctx->machine, ctx->op);
 
   return Qnil;
 }
@@ -307,343 +261,443 @@ VALUE um_timeout(struct um *machine, VALUE interval, VALUE class) {
   static ID ID_new = 0;
   if (!ID_new) ID_new = rb_intern("new");
 
-  struct um_op *op = um_op_checkout(machine);
+  struct um_op *op = um_op_alloc(machine);
+  um_prep_op(machine, op, OP_TIMEOUT);
   op->ts = um_double_to_timespec(NUM2DBL(interval));
+  RB_OBJ_WRITE(machine->self, &op->fiber, rb_fiber_current());
+  RB_OBJ_WRITE(machine->self, &op->value, rb_funcall(class, ID_new, 0));
 
   struct io_uring_sqe *sqe = um_get_sqe(machine, op);
   io_uring_prep_timeout(sqe, &op->ts, 0, 0);
-  op->state = OP_submitted;
-  op->fiber = rb_fiber_current();
-  op->resume_value = rb_funcall(class, ID_new, 0);
 
-  struct op_ensure_ctx ctx = { .machine = machine, .op = op };
+  struct op_ctx ctx = { .machine = machine, .op = op };
   return rb_ensure(rb_yield, Qnil, um_timeout_ensure, (VALUE)&ctx);
 }
 
-inline void discard_op_if_completed(struct um *machine, struct um_op *op) {
-  if (op->state == OP_completed) um_op_checkin(machine, op);
-}
+/*******************************************************************************
+                         blocking singleshot ops
+*******************************************************************************/
 
-inline VALUE um_sleep(struct um *machine, double duration) {
-  struct um_op *op = um_op_checkout(machine);
-  op->ts = um_double_to_timespec(duration);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  __s32 result = 0;
-
-  io_uring_prep_timeout(sqe, &op->ts, 0, 0);
-  op->state = OP_submitted;
+VALUE um_sleep(struct um *machine, double duration) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_SLEEP);
+  op.ts = um_double_to_timespec(duration);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_timeout(sqe, &op.ts, 0, 0);
+  VALUE ret = um_fiber_switch(machine);
 
-  um_await_op(machine, op, &result, NULL);
+  if (!um_op_completed_p(&op))
+    um_cancel_and_wait(machine, &op);
+  else {
+    if (op.result.res != -ETIME) um_raise_on_error_result(op.result.res);
+    ret = DBL2NUM(duration);
+  }
 
-  discard_op_if_completed(machine, op);
-  if (result != -ETIME) um_raise_on_system_error(result);
-  return Qnil;
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
 inline VALUE um_read(struct um *machine, int fd, VALUE buffer, int maxlen, int buffer_offset) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  __s32 result = 0;
-  __u32 flags = 0;
-
+  struct um_op op;
+  um_prep_op(machine, &op, OP_READ);
+  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);
-  op->state = OP_submitted;
+  
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op)) {
+    um_update_read_buffer(machine, buffer, buffer_offset, op.result.res, op.result.flags);
+    ret = INT2NUM(op.result.res);
+    
+  }
 
-  um_await_op(machine, op, &result, &flags);
+  RB_GC_GUARD(buffer);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
+}
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  um_update_read_buffer(machine, buffer, buffer_offset, result, flags);
-  return INT2FIX(result);
+VALUE um_write(struct um *machine, int fd, VALUE str, int len) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_WRITE);
+  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);
+
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
+
+  RB_GC_GUARD(str);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_multishot_ensure(VALUE arg) {
-  struct op_ensure_ctx *ctx = (struct op_ensure_ctx *)arg;
+VALUE um_close(struct um *machine, int fd) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_CLOSE);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_close(sqe, fd);
 
-  switch (ctx->op->state) {
-    case OP_submitted:
-      um_cancel_op(ctx->machine, ctx->op);
-      break;
-    case OP_completed:
-      um_op_checkin(ctx->machine, ctx->op);
-      break;
-    default:
-  }
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(fd);
 
-  if (ctx->read_buf) free(ctx->read_buf);
-  return Qnil;
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-static inline void um_read_each_prepare_op(struct op_ensure_ctx *ctx, int singleshot_mode) {
-  struct um_op *op = um_op_checkout(ctx->machine);
-  struct io_uring_sqe *sqe = um_get_sqe(ctx->machine, op);
+VALUE um_accept(struct um *machine, int fd) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_ACCEPT);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_accept(sqe, fd, NULL, NULL, 0);
 
-  if (singleshot_mode)
-    io_uring_prep_read(sqe, ctx->fd, ctx->read_buf, ctx->read_maxlen, -1);
-  else {
-    io_uring_prep_read_multishot(sqe, ctx->fd, 0, -1, ctx->bgid);
-    op->is_multishot = 1;
-  }
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
 
-  op->state = OP_submitted;
-  ctx->op = op;
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-int um_read_each_safe_loop_singleshot(struct op_ensure_ctx *ctx, int total) {
-  struct buf_ring_descriptor *desc = ctx->machine->buffer_rings + ctx->bgid;
-  __s32 result = 0;
-  ctx->read_maxlen = desc->buf_size;
-  ctx->read_buf = malloc(desc->buf_size);
+VALUE um_socket(struct um *machine, int domain, int type, int protocol, uint flags) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_SOCKET);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_socket(sqe, domain, type, protocol, flags);
 
-  while (1) {
-    um_read_each_prepare_op(ctx, 1);
-    um_await_op(ctx->machine, ctx->op, &result, NULL);
-    um_raise_on_system_error(result);
-    if (!result) return total;
-
-    total += result;
-    VALUE buf = rb_str_new(ctx->read_buf, result);
-    rb_yield(buf);
-    um_op_checkin(ctx->machine, ctx->op);
-  }
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
+
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
+VALUE um_connect(struct um *machine, int fd, const struct sockaddr *addr, socklen_t addrlen) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_CONNECT);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_connect(sqe, fd, addr, addrlen);
 
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
 
-int um_read_each_multishot_process_results(struct op_ensure_ctx *ctx, int *total) {
-  __s32 result = 0;
-  __u32 flags = 0;
-  __s32 bad_result = 0;
-  int eof = 0;
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
+}
 
-  while (um_op_result_shift(ctx->machine, ctx->op, &result, &flags)) {
-    if (result < 0) {
-      bad_result = result;
-      break;
-    }
-    if (result == 0) {
-      eof = 1;
-      break;
-    }
+VALUE um_send(struct um *machine, int fd, VALUE buffer, int len, int flags) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_SEND);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_send(sqe, fd, RSTRING_PTR(buffer), len, flags);
 
-    *total += result;
-    VALUE buf = um_get_string_from_buffer_ring(ctx->machine, ctx->bgid, result, flags);
-    rb_yield(buf);
-  }
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
 
-  if (ctx->op->state == OP_completed) {
-    um_op_checkin(ctx->machine, ctx->op);
+  RB_GC_GUARD(buffer);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
+}
 
-    // TTY devices might not support multishot reads:
-    // https://github.com/axboe/liburing/issues/1185. A workaround is to
-    // fallback to singleshot mode, using the first buffer in the buffer
-    // group.
-    if (!(flags & IORING_CQE_F_BUFFER)) {
-      *total = um_read_each_safe_loop_singleshot(ctx, *total);
-      return 0;
-    }
-    else
-      um_read_each_prepare_op(ctx, 0);
+VALUE um_recv(struct um *machine, int fd, VALUE buffer, int maxlen, int flags) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_RECV);
+  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);
+  if (um_check_completion(machine, &op)) {
+    um_update_read_buffer(machine, buffer, 0, op.result.res, op.result.flags);
+    ret = INT2NUM(op.result.res);
   }
-  if (bad_result)
-    um_raise_on_system_error(bad_result);
 
-  return eof ? 0 : 1;
+  RB_GC_GUARD(buffer);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_read_each_safe_loop(VALUE arg) {
-  struct op_ensure_ctx *ctx = (struct op_ensure_ctx *)arg;
-  int total = 0;
-
-  um_read_each_prepare_op(ctx, 0);
+VALUE um_bind(struct um *machine, int fd, struct sockaddr *addr, socklen_t addrlen) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_BIND);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_bind(sqe, fd, addr, addrlen);
 
-  while (1) {
-    um_await_op(ctx->machine, ctx->op, NULL, NULL);
-    if (!ctx->op->results_head)
-      rb_raise(rb_eRuntimeError, "no result found!\n");
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
 
-    if (!um_read_each_multishot_process_results(ctx, &total))
-      return INT2NUM(total);
-  }
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_read_each(struct um *machine, int fd, int bgid) {
-  struct op_ensure_ctx ctx = { .machine = machine, .fd = fd, .bgid = bgid, .read_buf = NULL };
-  return rb_ensure(um_read_each_safe_loop, (VALUE)&ctx, um_multishot_ensure, (VALUE)&ctx);
+VALUE um_listen(struct um *machine, int fd, int backlog) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_BIND);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_listen(sqe, fd, backlog);
+
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
+
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_write(struct um *machine, int fd, VALUE buffer, int len) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  __s32 result = 0;
-  __u32 flags = 0;
+VALUE um_getsockopt(struct um *machine, int fd, int level, int opt) {
+  VALUE ret = Qnil;
+  int value;
 
-  io_uring_prep_write(sqe, fd, RSTRING_PTR(buffer), len, -1);
-  op->state = OP_submitted;
+#ifdef HAVE_IO_URING_PREP_CMD_SOCK
+  struct um_op op;
+  um_prep_op(machine, &op, OP_GETSOCKOPT);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_cmd_sock(sqe, SOCKET_URING_OP_GETSOCKOPT, fd, level, opt, &value, sizeof(value));
 
-  um_await_op(machine, op, &result, &flags);
+  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
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_close(struct um *machine, int fd) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  __s32 result = 0;
-  __u32 flags = 0;
+VALUE um_setsockopt(struct um *machine, int fd, int level, int opt, int value) {
+  VALUE ret = Qnil;
 
-  io_uring_prep_close(sqe, fd);
-  op->state = OP_submitted;
+#ifdef HAVE_IO_URING_PREP_CMD_SOCK
+  struct um_op op;
+  um_prep_op(machine, &op, OP_GETSOCKOPT);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_cmd_sock(sqe, SOCKET_URING_OP_SETSOCKOPT, fd, level, opt, &value, sizeof(value));
 
-  um_await_op(machine, op, &result, &flags);
+  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
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(fd);
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
 }
 
-VALUE um_accept(struct um *machine, int fd) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  struct sockaddr addr;
-  socklen_t len;
-  __s32 result = 0;
-  __u32 flags = 0;
+VALUE um_open(struct um *machine, VALUE pathname, int flags, int mode) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_BIND);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_open(sqe, StringValueCStr(pathname), flags, mode);
+
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
+
+  RB_GC_GUARD(ret);
+  return raise_if_exception(ret);
+}
+
+VALUE um_waitpid(struct um *machine, int pid, int options) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_BIND);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+
+  siginfo_t infop;
+  io_uring_prep_waitid(sqe, P_PID, pid, &infop, options, 0);
 
-  io_uring_prep_accept(sqe, fd, &addr, &len, 0);
-  op->state = OP_submitted;
+  VALUE ret = um_fiber_switch(machine);
+  if (um_check_completion(machine, &op))
+    ret = INT2NUM(op.result.res);
 
-  um_await_op(machine, op, &result, &flags);
+  RB_GC_GUARD(ret);
+  raise_if_exception(ret);
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+  return rb_ary_new_from_args(2, INT2NUM(infop.si_pid), INT2NUM(infop.si_status));
 }
 
-VALUE um_accept_each_safe_loop(VALUE arg) {
-  struct op_ensure_ctx *ctx = (struct op_ensure_ctx *)arg;
-  __s32 result = 0;
-  __u32 flags = 0;
+/*******************************************************************************
+                            multishot ops
+*******************************************************************************/
 
-  while (1) {
-    um_await_op(ctx->machine, ctx->op, &result, &flags);
-    if (!ctx->op->results_head) {
-      // this shouldn't happen!
-      rb_raise(rb_eRuntimeError, "no result found for accept_each loop");
-    }
+VALUE accept_each_begin(VALUE arg) {
+  struct op_ctx *ctx = (struct op_ctx *)arg;
+  struct io_uring_sqe *sqe = um_get_sqe(ctx->machine, ctx->op);
+  io_uring_prep_multishot_accept(sqe, ctx->fd, NULL, NULL, 0);
+
+  while (true) {
+    VALUE ret = um_fiber_switch(ctx->machine);
+    if (!um_op_completed_p(ctx->op))
+      return raise_if_exception(ret);
 
-    while (um_op_result_shift(ctx->machine, ctx->op, &result, &flags)) {
-      um_raise_on_system_error(result);
-      if (likely(result > 0))
-        rb_yield(INT2FIX(result));
-      else
+    int more = false;
+    struct um_op_result *result = &ctx->op->result;
+    while (result) {
+      more = (result->flags & IORING_CQE_F_MORE);
+      if (result->res < 0) {
+        um_op_multishot_results_clear(ctx->machine, ctx->op);
         return Qnil;
+      }
+      rb_yield(INT2NUM(result->res));
+      result = result->next;
     }
+    um_op_multishot_results_clear(ctx->machine, ctx->op);
+    if (more)
+      ctx->op->flags &= ~OP_F_COMPLETED;
+    else
+      break;
   }
-}
-
-VALUE um_accept_each(struct um *machine, int fd) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  io_uring_prep_multishot_accept(sqe, fd, NULL, NULL, 0);
-  op->state = OP_submitted;
-  op->is_multishot = 1;
 
-  struct op_ensure_ctx ctx = { .machine = machine, .op = op, .read_buf = NULL };
-  return rb_ensure(um_accept_each_safe_loop, (VALUE)&ctx, um_multishot_ensure, (VALUE)&ctx);
+  return Qnil;
 }
 
-VALUE um_socket(struct um *machine, int domain, int type, int protocol, uint flags) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
+VALUE multishot_ensure(VALUE arg) {
+  struct op_ctx *ctx = (struct op_ctx *)arg;
+  if (ctx->op->multishot_result_count) {
+    int more = ctx->op->multishot_result_tail->flags & IORING_CQE_F_MORE;
+    if (more)
+      ctx->op->flags &= ~OP_F_COMPLETED;
+    um_op_multishot_results_clear(ctx->machine, ctx->op);
+  }
+  if (!um_op_completed_p(ctx->op))
+    um_cancel_and_wait(ctx->machine, ctx->op);
 
-  io_uring_prep_socket(sqe, domain, type, protocol, flags);
-  op->state = OP_submitted;
+  if (ctx->read_buf)
+    free(ctx->read_buf);
 
-  um_await_op(machine, op, &result, NULL);
+  return Qnil;
+}
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+VALUE um_accept_each(struct um *machine, int fd) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_ACCEPT_MULTISHOT);
+  
+  struct op_ctx ctx = { .machine = machine, .op = &op, .fd = fd, .read_buf = NULL };
+  return rb_ensure(accept_each_begin, (VALUE)&ctx, multishot_ensure, (VALUE)&ctx);
 }
 
-VALUE um_connect(struct um *machine, int fd, const struct sockaddr *addr, socklen_t addrlen) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
+int um_read_each_singleshot_loop(struct op_ctx *ctx) {
+  struct buf_ring_descriptor *desc = ctx->machine->buffer_rings + ctx->bgid;
+  ctx->read_maxlen = desc->buf_size;
+  ctx->read_buf = malloc(desc->buf_size);
+  int total = 0;
 
-  io_uring_prep_connect(sqe, fd, addr, addrlen);
-  op->state = OP_submitted;
+  while (1) {
+    um_prep_op(ctx->machine, ctx->op, OP_READ);
+    struct io_uring_sqe *sqe = um_get_sqe(ctx->machine, ctx->op);
+    io_uring_prep_read(sqe, ctx->fd, ctx->read_buf, ctx->read_maxlen, -1);
 
-  um_await_op(machine, op, &result, NULL);
+    VALUE ret = um_fiber_switch(ctx->machine);
+    if (um_op_completed_p(ctx->op)) {
+      um_raise_on_error_result(ctx->op->result.res);
+      if (!ctx->op->result.res) return total;
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+      VALUE buf = rb_str_new(ctx->read_buf, ctx->op->result.res);
+      total += ctx->op->result.res;
+      rb_yield(buf);
+      RB_GC_GUARD(buf);
+    }
+    else
+      return raise_if_exception(ret);
+  }
 }
 
-VALUE um_send(struct um *machine, int fd, VALUE buffer, int len, int flags) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
+// // returns true if more results are expected
+int read_recv_each_multishot_process_result(struct op_ctx *ctx, struct um_op_result *result, int *total) {
+  if (result->res == 0)
+    return false;
 
-  io_uring_prep_send(sqe, fd, RSTRING_PTR(buffer), len, flags);
-  op->state = OP_submitted;
+  *total += result->res;
+  VALUE buf = um_get_string_from_buffer_ring(ctx->machine, ctx->bgid, result->res, result->flags);
+  rb_yield(buf);
+  RB_GC_GUARD(buf);
 
-  um_await_op(machine, op, &result, NULL);
+  // TTY devices might not support multishot reads:
+  // https://github.com/axboe/liburing/issues/1185. We detect this by checking
+  // if the F_MORE flag is absent, then switch to single shot mode.
+  if (unlikely(!(result->flags & IORING_CQE_F_MORE))) {
+    *total += um_read_each_singleshot_loop(ctx);
+    return false;
+  }
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+  return true;
 }
 
-VALUE um_recv(struct um *machine, int fd, VALUE buffer, int maxlen, int flags) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
-
-  void *ptr = um_prepare_read_buffer(buffer, maxlen, 0);
-  io_uring_prep_recv(sqe, fd, ptr, maxlen, flags);
-  op->state = OP_submitted;
+void read_recv_each_prep(struct io_uring_sqe *sqe, struct op_ctx *ctx) {
+  switch (ctx->op->kind) {
+    case OP_READ_MULTISHOT:
+      io_uring_prep_read_multishot(sqe, ctx->fd, 0, -1, ctx->bgid);
+      return;
+    case OP_RECV_MULTISHOT:
+      io_uring_prep_recv_multishot(sqe, ctx->fd, NULL, 0, 0);
+	    sqe->buf_group = ctx->bgid;
+	    sqe->flags |= IOSQE_BUFFER_SELECT;
+      return;
+    default:
+      return;
+  }
+}
 
-  um_await_op(machine, op, &result, NULL);
+VALUE read_recv_each_begin(VALUE arg) {
+  struct op_ctx *ctx = (struct op_ctx *)arg;
+  struct io_uring_sqe *sqe = um_get_sqe(ctx->machine, ctx->op);
+  read_recv_each_prep(sqe, ctx);
+  int total = 0;
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  um_update_read_buffer(machine, buffer, 0, result, flags);
-  return INT2FIX(result);
-}
+  while (true) {
+    VALUE ret = um_fiber_switch(ctx->machine);
+    if (!um_op_completed_p(ctx->op))
+      return raise_if_exception(ret);
 
-VALUE um_bind(struct um *machine, int fd, struct sockaddr *addr, socklen_t addrlen) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
+    int more = false;
+    struct um_op_result *result = &ctx->op->result;
+    while (result) {
+      um_raise_on_error_result(result->res);
 
-  io_uring_prep_bind(sqe, fd, addr, addrlen);
-  op->state = OP_submitted;
+      more = (result->flags & IORING_CQE_F_MORE);
+      if (!read_recv_each_multishot_process_result(ctx, result, &total))
+        return Qnil;
 
-  um_await_op(machine, op, &result, NULL);
+      // rb_yield(INT2NUM(result->res));
+      result = result->next;
+    }
+    um_op_multishot_results_clear(ctx->machine, ctx->op);
+    if (more)
+      ctx->op->flags &= ~OP_F_COMPLETED;
+    else
+      break;
+  }
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
+  return Qnil;
 }
 
-VALUE um_listen(struct um *machine, int fd, int backlog) {
-  struct um_op *op = um_op_checkout(machine);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
-  int result = 0;
+VALUE um_read_each(struct um *machine, int fd, int bgid) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_READ_MULTISHOT);
 
-  io_uring_prep_listen(sqe, fd, backlog);
-  op->state = OP_submitted;
+  struct op_ctx ctx = { .machine = machine, .op = &op, .fd = fd, .bgid = bgid, .read_buf = NULL };
+  return rb_ensure(read_recv_each_begin, (VALUE)&ctx, multishot_ensure, (VALUE)&ctx);
+}
 
-  um_await_op(machine, op, &result, NULL);
+VALUE um_recv_each(struct um *machine, int fd, int bgid, int flags) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_RECV_MULTISHOT);
 
-  discard_op_if_completed(machine, op);
-  um_raise_on_system_error(result);
-  return INT2FIX(result);
-}
+  struct op_ctx ctx = { .machine = machine, .op = &op, .fd = fd, .bgid = bgid, .read_buf = NULL, .flags = flags };
+  return rb_ensure(read_recv_each_begin, (VALUE)&ctx, multishot_ensure, (VALUE)&ctx);
+}