uringmachine 0.25.0 → 0.26.0

data/ext/um/um.h

@@ -73,20 +73,31 @@ enum um_op_kind {
   OP_READ_MULTISHOT,
   OP_RECV_MULTISHOT,
   OP_TIMEOUT_MULTISHOT,
-  OP_SLEEP_MULTISHOT
 };
 
-#define OP_F_COMPLETED        (1U <<  0) // op is completed (set on each CQE for multishot ops)
-#define OP_F_TRANSIENT        (1U <<  1) // op is heap allocated
-#define OP_F_ASYNC            (1U <<  2) // op belongs to an AsyncOp
-#define OP_F_CANCELED         (1U <<  3) // op is cancelled
-#define OP_F_IGNORE_CANCELED  (1U <<  4) // CQE with -ECANCEL should be ignored
-#define OP_F_MULTISHOT        (1U <<  5) // op is multishot
-#define OP_F_FREE_ON_COMPLETE (1U <<  6) // op should be freed on receiving CQE
-#define OP_F_RUNQUEUE_SKIP    (1U <<  7) // runqueue entry should be skipped
-#define OP_F_SELECT_POLLIN    (1U <<  8) // select POLLIN
-#define OP_F_SELECT_POLLOUT   (1U <<  9) // select POLLOUT
-#define OP_F_SELECT_POLLPRI   (1U << 10) // select POLLPRI
+
+#define OP_F_CQE_SEEN       (1U << 0) // CQE has been seen
+#define OP_F_CQE_DONE       (1U << 1) // CQE has been seen and operation is done
+#define OP_F_SCHEDULED      (1U << 2) // op is on runqueue
+#define OP_F_CANCELED       (1U << 3) // op is cancelled (disregard CQE results)
+#define OP_F_MULTISHOT      (1U << 4) // op is multishot
+#define OP_F_ASYNC          (1U << 5) // op is async (no fiber is scheduled to be resumed on completion)
+#define OP_F_TRANSIENT      (1U << 6) // op is on transient list (for GC purposes)
+#define OP_F_FREE_IOVECS    (1U << 7) // op->iovecs should be freed on release
+#define OP_F_SKIP           (1U << 8) // op should be skipped when pulled from runqueue
+
+#define OP_F_SELECT_POLLIN  (1U << 13) // select POLLIN
+#define OP_F_SELECT_POLLOUT (1U << 14) // select POLLOUT
+#define OP_F_SELECT_POLLPRI (1U << 15) // select POLLPRI
+
+#define OP_CQE_SEEN_P(op)   ((op)->flags & OP_F_CQE_SEEN)
+#define OP_CQE_DONE_P(op)   ((op)->flags & OP_F_CQE_DONE)
+#define OP_SCHEDULED_P(op)  ((op)->flags & OP_F_SCHEDULED)
+#define OP_CANCELED_P(op)   ((op)->flags & OP_F_CANCELED)
+#define OP_MULTISHOT_P(op)  ((op)->flags & OP_F_MULTISHOT)
+#define OP_ASYNC_P(op)      ((op)->flags & OP_F_ASYNC)
+#define OP_TRANSIENT_P(op)  ((op)->flags & OP_F_TRANSIENT)
+#define OP_SKIP_P(op)       ((op)->flags & OP_F_SKIP)
 
 struct um_op_result {
   __s32 res;
@@ -100,6 +111,7 @@ struct um_op {
 
   enum um_op_kind kind;
   uint flags;
+  uint ref_count;
 
   VALUE fiber;
   VALUE value;
@@ -109,7 +121,12 @@ struct um_op {
   struct um_op_result *multishot_result_tail;
   uint multishot_result_count;
 
-  struct __kernel_timespec ts; // used for timeout operation
+	union {
+    struct __kernel_timespec ts; // used for timeout operation
+    struct iovec *iovecs; // used for vectorized write/send
+    siginfo_t siginfo; // used for waitid
+    int int_value; // used for getsockopt
+  };
 };
 
 struct um_buffer {
@@ -237,8 +254,12 @@ void um_teardown(struct um *machine);
 VALUE um_metrics(struct um *machine, struct um_metrics *metrics);
 
 const char * um_op_kind_name(enum um_op_kind kind);
-struct um_op *um_op_alloc(struct um *machine);
-void um_op_free(struct um *machine, struct um_op *op);
+
+struct um_op *um_op_acquire(struct um *machine);
+void um_op_release(struct um *machine, struct um_op *op);
+
+// struct um_op *um_op_alloc(struct um *machine);
+// void um_op_free(struct um *machine, struct um_op *op);
 void um_op_clear(struct um *machine, struct um_op *op);
 void um_op_transient_add(struct um *machine, struct um_op *op);
 void um_op_transient_remove(struct um *machine, struct um_op *op);
@@ -264,9 +285,9 @@ VALUE um_raise_exception(VALUE v);
 
 #define RAISE_IF_EXCEPTION(v) if (unlikely(um_value_is_exception_p(v))) { um_raise_exception(v); }
 
-void um_prep_op(struct um *machine, struct um_op *op, enum um_op_kind kind, unsigned flags);
+void um_prep_op(struct um *machine, struct um_op *op, enum um_op_kind kind, uint ref_count, unsigned flags);
 void um_raise_on_error_result(int result);
-void um_get_buffer_bytes_for_writing(VALUE buffer, const void **base, size_t *size);
+int um_get_buffer_bytes_for_writing(VALUE buffer, const void **base, size_t *size, int raise_on_bad_buffer);
 void * um_prepare_read_buffer(VALUE buffer, ssize_t len, ssize_t ofs);
 void um_update_read_buffer(VALUE buffer, ssize_t buffer_offset, __s32 result);
 int um_setup_buffer_ring(struct um *machine, unsigned size, unsigned count);
@@ -283,10 +304,9 @@ VALUE um_yield(struct um *machine);
 VALUE um_switch(struct um *machine);
 VALUE um_wakeup(struct um *machine);
 void um_cancel_op(struct um *machine, struct um_op *op);
-void um_cancel_and_wait(struct um *machine, struct um_op *op);
-int um_check_completion(struct um *machine, struct um_op *op);
-
-#define um_op_completed_p(op) ((op)->flags & OP_F_COMPLETED)
+void um_cancel_op_and_discard_cqe(struct um *machine, struct um_op *op);
+void um_cancel_op_and_await_cqe(struct um *machine, struct um_op *op);
+int um_verify_op_completion(struct um *machine, struct um_op *op, int await_cancelled);
 
 void um_schedule(struct um *machine, VALUE fiber, VALUE value);
 VALUE um_timeout(struct um *machine, VALUE interval, VALUE class);

data/ext/um/um_async_op.c

@@ -5,11 +5,11 @@ VALUE um_prep_timeout(struct um *machine, double interval) {
   static ID ID_new = 0;
   if (!ID_new) ID_new = rb_intern("new");
 
-  struct um_op *op = malloc(sizeof(struct um_op));
-  um_prep_op(machine, op, OP_TIMEOUT, OP_F_TRANSIENT | OP_F_ASYNC);
+  struct um_op *op = um_op_acquire(machine);
+  um_prep_op(machine, op, OP_TIMEOUT, 2, OP_F_ASYNC | OP_F_TRANSIENT);
   op->ts = um_double_to_timespec(interval);
 
-  VALUE obj = rb_funcall(cAsyncOp, rb_intern_const("new"), 0);
+  VALUE obj = rb_funcall(cAsyncOp, ID_new, 0);
   um_async_op_set(obj, machine, op);
 
   RB_OBJ_WRITE(machine->self, &op->async_op, obj);
@@ -17,18 +17,19 @@ VALUE um_prep_timeout(struct um *machine, double interval) {
   struct io_uring_sqe *sqe = um_get_sqe(machine, op);
   io_uring_prep_timeout(sqe, &op->ts, 0, 0);
 
-  um_op_transient_add(machine, op);
-
   return obj;
 }
 
 VALUE um_async_op_await(struct um_async_op *async_op) {
+  if (OP_CQE_DONE_P(async_op->op)) return async_op->op->value;
+
   RB_OBJ_WRITE(async_op->machine->self, &async_op->op->fiber, rb_fiber_current());
-  async_op->op->flags &= ~OP_F_ASYNC;
+  um_op_transient_remove(async_op->machine, async_op->op);
+  async_op->op->flags &= ~(OP_F_ASYNC | OP_F_TRANSIENT);
 
   VALUE ret = um_switch(async_op->machine);
-  if (!um_op_completed_p(async_op->op))
-    um_cancel_and_wait(async_op->machine, async_op->op);
+
+  um_verify_op_completion(async_op->machine, async_op->op, true);
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);

data/ext/um/um_async_op_class.c

@@ -19,8 +19,8 @@ static void AsyncOp_mark(void *ptr) {
 
 static void AsyncOp_free(void *ptr) {
   struct um_async_op *async_op = ptr;
-  if (async_op->op)
-    um_op_free(async_op->machine, async_op->op);
+  if (likely(async_op->op))
+    um_op_release(async_op->machine, async_op->op);
 }
 
 static const rb_data_type_t AsyncOp_type = {
@@ -63,7 +63,7 @@ inline void raise_on_missing_op(struct um_async_op *async_op) {
 }
 
 inline int async_op_is_done(struct um_async_op *async_op) {
-  return (async_op->op->flags & OP_F_COMPLETED);
+  return OP_CQE_DONE_P(async_op->op);
 }
 
 /* Returns the kind of asynchronous operation.

data/ext/um/um_op.c

@@ -37,7 +37,6 @@ const char * um_op_kind_name(enum um_op_kind kind) {
     case OP_READ_MULTISHOT:     return "OP_READ_MULTISHOT";
     case OP_RECV_MULTISHOT:     return "OP_RECV_MULTISHOT";
     case OP_TIMEOUT_MULTISHOT:  return "OP_TIMEOUT_MULTISHOT";
-    case OP_SLEEP_MULTISHOT:    return "OP_SLEEP_MULTISHOT";
     default:                    return "UNKNOWN_OP_KIND";
   }
 }
@@ -59,6 +58,7 @@ inline void um_op_transient_add(struct um *machine, struct um_op *op) {
 }
 
 inline void um_op_transient_remove(struct um *machine, struct um_op *op) {
+  op->flags &= ~OP_F_TRANSIENT;
   if (op->prev)
     op->prev->next = op->next;
   if (op->next)
@@ -185,3 +185,17 @@ inline void um_op_free(struct um *machine, struct um_op *op) {
   machine->op_freelist = op;
   machine->metrics.ops_free++;
 }
+
+inline struct um_op *um_op_acquire(struct um *machine) {
+  return um_op_alloc(machine);
+}
+
+inline void um_op_release(struct um *machine, struct um_op *op) {
+  op->ref_count--;
+  if (op->ref_count) return;
+
+  if (op->flags & OP_F_FREE_IOVECS) free(op->iovecs);
+  if (op->flags & OP_F_MULTISHOT)
+    um_op_multishot_results_clear(machine, op);
+  um_op_free(machine, op);
+}

data/ext/um/um_ssl.c

@@ -97,7 +97,7 @@ int um_ssl_write(struct um *machine, VALUE ssl_obj, VALUE buf, int len) {
   SSL *ssl = RTYPEDDATA_GET_DATA(ssl_obj);
   const void *base;
   size_t size;
-  um_get_buffer_bytes_for_writing(buf, &base, &size);
+  um_get_buffer_bytes_for_writing(buf, &base, &size, true);
   if ((len == (int)-1) || (len > (int)size)) len = (int)size;
   if (unlikely(!len)) return INT2NUM(0);
 

data/ext/um/um_sync.c

@@ -6,35 +6,42 @@
 
 // The value argument is the current (known) futex value.
 void um_futex_wait(struct um *machine, uint32_t *futex, uint32_t value) {
-  struct um_op op;
-  um_prep_op(machine, &op, OP_FUTEX_WAIT, 0);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  struct um_op *op = um_op_acquire(machine);
+  um_prep_op(machine, op, OP_FUTEX_WAIT, 2, 0);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
   io_uring_prep_futex_wait(
     sqe, (uint32_t *)futex, value, FUTEX_BITSET_MATCH_ANY, FUTEX2_SIZE_U32, 0
   );
 
   VALUE ret = um_yield(machine);
-  if (!um_op_completed_p(&op))
-    um_cancel_and_wait(machine, &op);
+
+  if (unlikely(!OP_CQE_DONE_P(op)))
+    um_cancel_op_and_await_cqe(machine, op);
   else {
-    if (op.result.res != -EAGAIN)
-      um_raise_on_error_result(op.result.res);
+    int res = op->result.res;
+    if (res != -EAGAIN) {
+      um_op_release(machine, op);
+      um_raise_on_error_result(res);
+    }
   }
+  um_op_release(machine, op);
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);
 }
 
 void um_futex_wake(struct um *machine, uint32_t *futex, uint32_t num_waiters) {
-  struct um_op op;
-  um_prep_op(machine, &op, OP_FUTEX_WAKE, 0);
-  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  struct um_op *op = um_op_acquire(machine);
+  um_prep_op(machine, op, OP_FUTEX_WAKE, 2, 0);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, op);
   io_uring_prep_futex_wake(
     sqe, (uint32_t *)futex, num_waiters, FUTEX_BITSET_MATCH_ANY, FUTEX2_SIZE_U32, 0
   );
 
   VALUE ret = um_yield(machine);
-  um_check_completion(machine, &op);
+
+  um_verify_op_completion(machine, op, true);
+  um_op_release(machine, op);
 
   RAISE_IF_EXCEPTION(ret);
   RB_GC_GUARD(ret);

data/ext/um/um_utils.c

@@ -98,15 +98,21 @@ inline void um_update_read_buffer(VALUE buffer, ssize_t buffer_offset, __s32 res
   adjust_read_buffer_len(buffer, result, buffer_offset);
 }
 
-inline void um_get_buffer_bytes_for_writing(VALUE buffer, const void **base, size_t *size) {
+// returns false if buffer is invalid and raise_on_bad_buffer is false
+inline int um_get_buffer_bytes_for_writing(VALUE buffer, const void **base, size_t *size, int raise_on_bad_buffer) {
   if (TYPE(buffer) == T_STRING) {
     *base = RSTRING_PTR(buffer);
     *size = RSTRING_LEN(buffer);
   }
   else if (IO_BUFFER_P(buffer))
     rb_io_buffer_get_bytes_for_reading(buffer, base, size); // reading *from* buffer
-  else
-    um_raise_internal_error("Invalid buffer provided");
+  else {
+    if (raise_on_bad_buffer)
+      um_raise_internal_error("Invalid buffer provided");
+    else
+      return false;
+  }
+  return true;
 }
 
 int um_setup_buffer_ring(struct um *machine, unsigned size, unsigned count) {
@@ -204,7 +210,11 @@ inline struct iovec *um_alloc_iovecs_for_writing(int argc, VALUE *argv, size_t *
   size_t len = 0;
 
   for (int i = 0; i < argc; i++) {
-    um_get_buffer_bytes_for_writing(argv[i], (const void **)&iovecs[i].iov_base, &iovecs[i].iov_len);
+    int ok = um_get_buffer_bytes_for_writing(argv[i], (const void **)&iovecs[i].iov_base, &iovecs[i].iov_len, false);
+    if (unlikely(!ok)) {
+      free(iovecs);
+      um_raise_internal_error("Invalid buffer provided");
+    }
     len += iovecs[i].iov_len;
   }
   if (total_len) *total_len = len;

data/lib/uringmachine/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 class UringMachine
-  VERSION = '0.25.0'
+  VERSION = '0.26.0'
 end

data/test/test_async_op.rb

@@ -51,11 +51,12 @@ class AsyncOpTest < UMBaseTest
       @op.cancel
     end
 
-    res = @op.await
+    assert_raises(Errno::ECANCELED) {
+      @op.await
+    }
 
     assert_equal 0, machine.metrics[:ops_pending]
     assert_equal true, @op.done?
-    assert_equal (-ECANCELED), res
     assert_equal true, @op.cancelled?
   end
 

data/test/test_fiber_scheduler.rb

@@ -393,9 +393,12 @@ class FiberSchedulerTest < UMBaseTest
       sleep 0.001
       buf = IO.read(fn)
     end
-    assert_equal 2, machine.metrics[:total_ops]
 
     @scheduler.join
+    metrics = machine.metrics
+    assert_in_range 5..12, metrics[:total_ops]
+    assert_equal 0, metrics[:ops_pending]
+    assert_equal 256, metrics[:ops_free]
     assert_equal 'foobar', buf
     assert_equal({
       fiber: 2,