uringmachine 0.3 → 0.5

data/ext/um/um_op.c

@@ -1,151 +1,123 @@
 #include "um.h"
 
-inline struct um_result_entry *um_result_checkout(struct um *machine) {
-  if (machine->result_freelist) {
-    struct um_result_entry *entry = machine->result_freelist;
-    machine->result_freelist = entry->next;
-    return entry;
-  }
-
-  struct um_result_entry *entry = malloc(sizeof(struct um_result_entry));
-  return entry;
-}
-
-inline void um_result_checkin(struct um *machine, struct um_result_entry *entry) {
-  entry->next = machine->result_freelist;
-  machine->result_freelist = entry;  
+inline void um_op_clear(struct um *machine, struct um_op *op) {
+  memset(op, 0, sizeof(struct um_op));
+  RB_OBJ_WRITE(machine->self, &op->fiber, Qnil);
+  RB_OBJ_WRITE(machine->self, &op->value, Qnil);
 }
 
-inline void um_op_result_cleanup(struct um *machine, struct um_op *op) {  
-  struct um_result_entry *entry = op->results_head;
-  while (entry) {
-    struct um_result_entry *next = entry->next;
-    um_result_checkin(machine, entry);
-    entry = next;
+inline void um_op_transient_add(struct um *machine, struct um_op *op) {
+  if (machine->transient_head) {
+    op->next = machine->transient_head;
+    machine->transient_head->prev = op;
   }
-  op->results_head = op->results_tail = NULL;
+  machine->transient_head = op;
 }
 
-inline void um_op_result_push(struct um *machine, struct um_op *op, int result, int flags) {
-  struct um_result_entry *entry = um_result_checkout(machine);
-  entry->next = 0;
-  entry->result = result;
-  entry->flags = flags;
-  if (op->results_tail) {
-    op->results_tail->next = entry;
-    op->results_tail = entry;
-  }
-  else {
-    op->results_head = op->results_tail = entry;
-  }
-}
+inline void um_op_transient_remove(struct um *machine, struct um_op *op) {
+  if (op->prev)
+    op->prev->next = op->next;
+  if (op->next)
+    op->next->prev = op->prev;
 
-inline int um_op_result_shift(struct um *machine, struct um_op *op, int *result, int *flags) {
-  if (!op->results_head) return 0;
-
-  struct um_result_entry *entry = op->results_head;
-  *result = entry->result;
-  *flags = entry->flags;
-  op->results_head = entry->next;
-  if (!op->results_head)
-    op->results_tail = NULL;
-  um_result_checkin(machine, entry);
-  return 1;
+  if (machine->transient_head == op)
+    machine->transient_head = op->next;
 }
 
-inline void um_op_clear(struct um_op *op) {
-  memset(op, 0, sizeof(struct um_op));
-  op->fiber = op->resume_value = Qnil;
+inline void um_runqueue_push(struct um *machine, struct um_op *op) {
+  if (machine->runqueue_tail) {
+    op->prev = machine->runqueue_tail;
+    machine->runqueue_tail->next = op;
+    machine->runqueue_tail = op;
+  }
+  else
+    machine->runqueue_head = machine->runqueue_tail = op;
+  op->next = NULL;
 }
 
-inline struct um_op *um_op_checkout(struct um *machine) {
-  machine->pending_count++;
-
-  if (machine->op_freelist) {
-    struct um_op *op = machine->op_freelist;
-    machine->op_freelist = op->next;
-    um_op_clear(op);
-    return op;
-  }
+inline struct um_op *um_runqueue_shift(struct um *machine) {
+  struct um_op *op = machine->runqueue_head;
+  if (!op) return NULL;
 
-  struct um_op *op = malloc(sizeof(struct um_op));
-  um_op_clear(op);
+  machine->runqueue_head = op->next;
+  if (!machine->runqueue_head)
+    machine->runqueue_tail = NULL;
   return op;
 }
 
-inline void um_op_checkin(struct um *machine, struct um_op *op) {
-  um_op_result_cleanup(machine, op);
-
-  machine->pending_count--;
-
-  op->next = machine->op_freelist;
-  machine->op_freelist = op;
+inline void um_op_list_mark(struct um *machine, struct um_op *head) {
+  while (head) {
+    struct um_op *next = head->next;
+    rb_gc_mark_movable(head->fiber);
+    rb_gc_mark_movable(head->value);
+    head = next;
+  }
 }
 
-inline struct um_op *um_runqueue_find_by_fiber(struct um *machine, VALUE fiber) {
-  struct um_op *op = machine->runqueue_head;
-  while (op) {
-    if (op->fiber == fiber) return op;
-
-    op = op->next;
+inline void um_op_list_compact(struct um *machine, struct um_op *head) {
+  while (head) {
+    struct um_op *next = head->next;
+    head->fiber = rb_gc_location(head->fiber);
+    head->value = rb_gc_location(head->value);
+    head = next;
   }
-  return NULL;
 }
 
-inline void um_runqueue_push(struct um *machine, struct um_op *op) {
-  if (machine->runqueue_tail) {
-    op->prev = machine->runqueue_tail;
-    machine->runqueue_tail->next = op;
-    machine->runqueue_tail = op;
-  }
-  else {
-    op->prev = NULL;
-    machine->runqueue_head = machine->runqueue_tail = op;
+inline struct um_op_result *multishot_result_alloc(struct um *machine) {  
+  if (machine->result_freelist) {
+    struct um_op_result *result = machine->result_freelist;
+    machine->result_freelist = result->next;
+    return result;
   }
-  op->next = NULL;
+  return malloc(sizeof(struct um_op_result));
+}
+
+inline void multishot_result_free(struct um *machine, struct um_op_result *result) {
+  result->next = machine->result_freelist;
+  machine->result_freelist = result;
 }
 
-inline void um_runqueue_unshift(struct um *machine, struct um_op *op) {
-  if (machine->runqueue_head) {
-    op->next = machine->runqueue_head;
-    machine->runqueue_head->prev = op;
-    machine->runqueue_head = op;
+inline void um_op_multishot_results_push(struct um *machine, struct um_op *op, __s32 res, __u32 flags) {
+  if (!op->multishot_result_count) {
+    op->result.res    = res;
+    op->result.flags  = flags;
+    op->result.next   = NULL;
+    op->multishot_result_tail = &op->result;
   }
   else {
-    op->next = NULL;
-    machine->runqueue_head = machine->runqueue_tail = op;
+    struct um_op_result *result = multishot_result_alloc(machine);
+    result->res   = res;
+    result->flags = flags;
+    result->next  = NULL;
+    op->multishot_result_tail->next = result;
+    op->multishot_result_tail = result;
   }
-  op->prev = NULL;
+  op->multishot_result_count++;
 }
 
-inline struct um_op *um_runqueue_shift(struct um *machine) {
-  struct um_op *op = machine->runqueue_head;
-  if (!op) return NULL;
+inline void um_op_multishot_results_clear(struct um *machine, struct um_op *op) {
+  if (op->multishot_result_count < 1) return;
 
-  op->prev = NULL;
-  if (!op->next) {
-    machine->runqueue_head = machine->runqueue_tail = NULL;    
-  }
-  else {
-    machine->runqueue_head = op->next;
-    op->next = NULL;
+  struct um_op_result *result = op->result.next;
+  while (result) {
+    struct um_op_result *next = result->next;
+    multishot_result_free(machine, result);
+    result = next;
   }
-  return op;
+  op->multishot_result_tail = NULL;
+  op->multishot_result_count = 0;
 }
 
-inline void um_free_op_linked_list(struct um *machine, struct um_op *op) {
-  while (op) {
-    struct um_op *next = op->next;
-    um_op_result_cleanup(machine, op);
-    free(op);
-    op = next;
+inline struct um_op *um_op_alloc(struct um *machine) {
+  if (machine->op_freelist) {
+    struct um_op *op = machine->op_freelist;
+    machine->op_freelist = op->next;
+    return op;
   }
+  return malloc(sizeof(struct um_op));
 }
 
-inline void um_free_result_linked_list(struct um *machine, struct um_result_entry *entry) {
-  while (entry) {
-    struct um_result_entry *next = entry->next;
-    free(entry);
-    entry = next;
-  }
+inline void um_op_free(struct um *machine, struct um_op *op) {
+  op->next = machine->op_freelist;
+  machine->op_freelist = op;
 }

data/ext/um/um_queue_class.c

@@ -0,0 +1,58 @@
+#include "um.h"
+#include <stdlib.h>
+
+VALUE cQueue;
+
+static void Queue_mark(void *ptr) {
+  struct um_queue *queue = ptr;
+  um_queue_mark(queue);
+}
+
+static void Queue_compact(void *ptr) {
+  struct um_queue *queue = ptr;
+  um_queue_compact(queue);
+}
+
+static void Queue_free(void *ptr) {
+  struct um_queue *queue = ptr;
+  um_queue_free(queue);
+}
+
+static size_t Queue_size(const void *ptr) {
+  return sizeof(struct um_queue);
+}
+
+static const rb_data_type_t Queue_type = {
+    "UringMachineQueue",
+    {Queue_mark, Queue_free, Queue_size, Queue_compact},
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
+};
+
+static VALUE Queue_allocate(VALUE klass) {
+  struct um_queue *queue = malloc(sizeof(struct um_queue));
+  return TypedData_Wrap_Struct(klass, &Queue_type, queue);
+}
+
+inline struct um_queue *Queue_data(VALUE self) {
+  return RTYPEDDATA_DATA(self);
+}
+
+VALUE Queue_initialize(VALUE self) {
+  struct um_queue *queue = Queue_data(self);
+  RB_OBJ_WRITE(self, &queue->self, self);
+  um_queue_init(queue);
+  return self;
+}
+
+VALUE Queue_count(VALUE self) {
+  struct um_queue *queue = Queue_data(self);
+  return UINT2NUM(queue->count);
+}
+
+void Init_Queue(void) {
+  cQueue = rb_define_class_under(cUM, "Queue", rb_cObject);
+  rb_define_alloc_func(cQueue, Queue_allocate);
+
+  rb_define_method(cQueue, "initialize", Queue_initialize, 0);
+  rb_define_method(cQueue, "count", Queue_count, 0);
+}

data/ext/um/um_sync.c

@@ -0,0 +1,273 @@
+#include "um.h"
+#include <stdatomic.h>
+#include <linux/futex.h>
+
+#define FUTEX2_SIZE_U32		0x02
+
+void um_futex_wait(struct um *machine, uint32_t *futex, uint32_t expect) {
+  struct um_op op;
+  um_prep_op(machine, &op, OP_FUTEX_WAIT);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  io_uring_prep_futex_wait(
+    sqe, (uint32_t *)futex, expect, FUTEX_BITSET_MATCH_ANY,
+		FUTEX2_SIZE_U32, 0
+  );
+  
+  VALUE ret = um_fiber_switch(machine);
+  if (!um_op_completed_p(&op))
+    um_cancel_and_wait(machine, &op);
+  else {
+    if (op.result.res != -EAGAIN)
+      um_raise_on_error_result(op.result.res);
+  }
+
+  RB_GC_GUARD(ret);
+  raise_if_exception(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);
+  struct io_uring_sqe *sqe = um_get_sqe(machine, &op);
+  // submit futex_wait
+  io_uring_prep_futex_wake(
+    sqe, (uint32_t *)futex, num_waiters, FUTEX_BITSET_MATCH_ANY,
+		FUTEX2_SIZE_U32, 0
+  );
+
+  VALUE ret = um_fiber_switch(machine);
+  um_check_completion(machine, &op);
+
+  RB_GC_GUARD(ret);
+  raise_if_exception(ret);
+}
+
+void um_futex_wake_transient(struct um *machine, uint32_t *futex, uint32_t num_waiters) {
+  struct io_uring_sqe *sqe = um_get_sqe(machine, NULL);
+  io_uring_prep_futex_wake(
+    sqe, (uint32_t *)futex, num_waiters, FUTEX_BITSET_MATCH_ANY,
+		FUTEX2_SIZE_U32, 0
+  );
+}
+
+
+#define MUTEX_LOCKED    1
+#define MUTEX_UNLOCKED  0
+
+void um_mutex_init(struct um_mutex *mutex) {
+  mutex->state = MUTEX_UNLOCKED;
+}
+
+inline void um_mutex_lock(struct um *machine, uint32_t *state) {
+  while (*state == MUTEX_LOCKED) {
+    um_futex_wait(machine, state, MUTEX_LOCKED);
+  }
+  *state = MUTEX_LOCKED;
+}
+
+inline void um_mutex_unlock(struct um *machine, uint32_t *state) {
+  *state = MUTEX_UNLOCKED;
+  // Wake up 1 waiting fiber
+  um_futex_wake(machine, state, 1);
+}
+
+struct sync_ctx {
+  struct um *machine;
+  uint32_t *state;
+};
+
+VALUE synchronize_begin(VALUE arg) {
+  struct sync_ctx *ctx = (struct sync_ctx *)arg;
+  um_mutex_lock(ctx->machine, ctx->state);
+  return rb_yield(Qnil);
+}
+
+VALUE synchronize_ensure(VALUE arg) {
+  struct sync_ctx *ctx = (struct sync_ctx *)arg;
+  um_mutex_unlock(ctx->machine, ctx->state);
+  return Qnil;
+}
+
+inline VALUE um_mutex_synchronize(struct um *machine, uint32_t *state) {
+  struct sync_ctx ctx = { .machine = machine, .state = state };
+  return rb_ensure(synchronize_begin, (VALUE)&ctx, synchronize_ensure, (VALUE)&ctx);
+}
+
+#define QUEUE_EMPTY 0
+#define QUEUE_READY 1
+
+inline void um_queue_init(struct um_queue *queue) {
+  queue->head = queue->tail = queue->free_head = NULL;
+  queue->state = QUEUE_EMPTY;
+  queue->count = 0;
+}
+
+inline void um_queue_free(struct um_queue *queue) {
+  struct um_queue_entry *entry = queue->head;
+  while (entry) {
+    struct um_queue_entry *next = entry->next;
+    free(entry);
+    entry = next;
+  }
+
+  entry = queue->free_head;
+  while (entry) {
+    struct um_queue_entry *next = entry->next;
+    free(entry);
+    entry = next;
+  }
+
+  free(queue);
+}
+
+inline void um_queue_mark(struct um_queue *queue) {
+  rb_gc_mark_movable(queue->self);
+  struct um_queue_entry *entry = queue->head;
+  while (entry) {
+    rb_gc_mark_movable(entry->value);
+    entry = entry->next;
+  }
+}
+
+inline void um_queue_compact(struct um_queue *queue) {
+  queue->self = rb_gc_location(queue->self);
+  struct um_queue_entry *entry = queue->head;
+  while (entry) {
+    entry->value = rb_gc_location(entry->value);
+    entry = entry->next;
+  }
+}
+
+inline struct um_queue_entry *um_queue_entry_checkout(struct um_queue *queue) {
+  struct um_queue_entry *entry = queue->free_head;
+  if (entry) {
+    queue->free_head = entry->next;
+  }
+  else
+    entry = malloc(sizeof(struct um_queue_entry));
+  return entry;
+}
+
+inline void um_queue_entry_checkin(struct um_queue *queue, struct um_queue_entry *entry) {
+  entry->next = queue->free_head;
+  queue->free_head = entry;
+}
+
+static inline void queue_add_head(struct um_queue *queue, VALUE value) {
+  struct um_queue_entry *entry = um_queue_entry_checkout(queue);
+
+  entry->next = queue->head;
+  if (queue->head) { 
+    queue->head->prev = entry;
+    queue->head = entry;
+  }
+  else
+    queue->head = queue->tail = entry;
+  entry->prev = NULL;
+  RB_OBJ_WRITE(queue->self, &entry->value, value);
+}
+
+static inline void queue_add_tail(struct um_queue *queue, VALUE value) {
+  struct um_queue_entry *entry = um_queue_entry_checkout(queue);
+
+  entry->prev = queue->tail;
+  if (queue->tail) { 
+    queue->tail->next = entry;
+    queue->tail = entry;
+  }
+  else
+    queue->head = queue->tail = entry;
+  entry->next = NULL;
+  RB_OBJ_WRITE(queue->self, &entry->value, value);
+}
+
+VALUE queue_remove_head(struct um_queue *queue) {
+  struct um_queue_entry *entry = queue->head;
+  queue->head = entry->next;
+  if (!queue->head) queue->tail = NULL;
+
+  VALUE v = entry->value;
+  um_queue_entry_checkin(queue, entry);
+  return v;
+
+}
+
+VALUE queue_remove_tail(struct um_queue *queue) {
+  struct um_queue_entry *entry = queue->tail;
+  queue->tail = entry->prev;
+  if (!queue->tail) queue->head = NULL;
+
+  VALUE v = entry->value;
+  um_queue_entry_checkin(queue, entry);
+  return v;
+}
+
+static inline VALUE um_queue_add(struct um *machine, struct um_queue *queue, VALUE value, int add_head) {
+  if (add_head) queue_add_head(queue, value);
+  else          queue_add_tail(queue, value);
+  
+  queue->count++;
+
+  queue->state = QUEUE_READY;
+  if (queue->num_waiters)
+    um_futex_wake_transient(machine, &queue->state, 1);
+  return queue->self;
+}
+
+VALUE um_queue_push(struct um *machine, struct um_queue *queue, VALUE value) {
+  return um_queue_add(machine, queue, value, false);
+}
+
+VALUE um_queue_unshift(struct um *machine, struct um_queue *queue, VALUE value) {
+  return um_queue_add(machine, queue, value, true);
+}
+
+enum queue_op { QUEUE_POP, QUEUE_SHIFT };
+
+struct queue_wait_ctx {
+  struct um *machine;
+  struct um_queue *queue;
+  enum queue_op op;
+};
+
+VALUE um_queue_remove_begin(VALUE arg) {
+  struct queue_wait_ctx *ctx = (struct queue_wait_ctx *)arg;
+  
+  ctx->queue->num_waiters++;
+  while (ctx->queue->state == QUEUE_EMPTY) {
+    um_futex_wait(ctx->machine, &ctx->queue->state, QUEUE_EMPTY);
+  }
+
+  if (ctx->queue->state != QUEUE_READY)
+    rb_raise(rb_eRuntimeError, "Internal error: queue should be in ready state!");
+  if (!ctx->queue->tail)
+    rb_raise(rb_eRuntimeError, "Internal error: queue should be in ready state!");
+
+  ctx->queue->count--;
+  return (ctx->op == QUEUE_POP ? queue_remove_tail : queue_remove_head)(ctx->queue);
+}
+
+VALUE um_queue_remove_ensure(VALUE arg) {
+  struct queue_wait_ctx *ctx = (struct queue_wait_ctx *)arg;
+
+  ctx->queue->num_waiters--;
+
+  if (ctx->queue->num_waiters && ctx->queue->tail) {
+    um_futex_wake_transient(ctx->machine, &ctx->queue->state, 1);
+  }
+  else if (!ctx->queue->tail) {
+    ctx->queue->state = QUEUE_EMPTY;
+  }
+
+  return Qnil;
+}
+
+VALUE um_queue_pop(struct um *machine, struct um_queue *queue) {
+  struct queue_wait_ctx ctx = { .machine = machine, .queue = queue, .op = QUEUE_POP };
+  return rb_ensure(um_queue_remove_begin, (VALUE)&ctx, um_queue_remove_ensure, (VALUE)&ctx);
+}
+
+VALUE um_queue_shift(struct um *machine, struct um_queue *queue) {
+  struct queue_wait_ctx ctx = { .machine = machine, .queue = queue, .op = QUEUE_SHIFT };
+  return rb_ensure(um_queue_remove_begin, (VALUE)&ctx, um_queue_remove_ensure, (VALUE)&ctx);
+}

data/ext/um/um_utils.c

@@ -1,4 +1,6 @@
 #include "um.h"
+#include <sys/mman.h>
+#include <stdlib.h>
 
 inline struct __kernel_timespec um_double_to_timespec(double value) {
   double integral;
@@ -22,7 +24,7 @@ inline VALUE um_raise_exception(VALUE e) {
   return rb_funcall(rb_mKernel, ID_raise, 1, e);
 }
 
-inline void um_raise_on_system_error(int result) {
+inline void um_raise_on_error_result(int result) {
   if (unlikely(result < 0)) rb_syserr_fail(-result, strerror(-result));
 }
 
@@ -46,13 +48,56 @@ static inline void adjust_read_buffer_len(VALUE buffer, int result, int ofs) {
   rb_str_set_len(buffer, len + (unsigned)ofs);
 }
 
-inline void um_update_read_buffer(struct um *machine, VALUE buffer, int buffer_offset, int result, int flags) {
+inline void um_update_read_buffer(struct um *machine, VALUE buffer, int buffer_offset, __s32 result, __u32 flags) {
   if (!result) return;
 
   adjust_read_buffer_len(buffer, result, buffer_offset);
 }
 
-inline VALUE get_string_from_buffer_ring(struct um *machine, int bgid, int result, int flags) {
+int um_setup_buffer_ring(struct um *machine, unsigned size, unsigned count) {
+  if (machine->buffer_ring_count == BUFFER_RING_MAX_COUNT)
+    rb_raise(rb_eRuntimeError, "Cannot setup more than BUFFER_RING_MAX_COUNT buffer rings");
+
+  struct buf_ring_descriptor *desc = machine->buffer_rings + machine->buffer_ring_count;
+  desc->buf_count = count;
+  desc->buf_size = size;
+
+  desc->br_size = sizeof(struct io_uring_buf) * desc->buf_count;
+	void *mapped = mmap(
+    NULL, desc->br_size, PROT_READ | PROT_WRITE,
+		MAP_ANONYMOUS | MAP_PRIVATE, 0, 0
+  );
+  if (mapped == MAP_FAILED)
+    rb_raise(rb_eRuntimeError, "Failed to allocate buffer ring");
+
+  desc->br = (struct io_uring_buf_ring *)mapped;
+  io_uring_buf_ring_init(desc->br);
+
+  unsigned bg_id = machine->buffer_ring_count;
+  int ret;
+  desc->br = io_uring_setup_buf_ring(&machine->ring, count, bg_id, 0, &ret);
+  if (!desc->br) {
+    munmap(desc->br, desc->br_size);
+    rb_syserr_fail(ret, strerror(ret));
+  }
+
+  if (posix_memalign(&desc->buf_base, 4096, desc->buf_count * desc->buf_size)) {
+    io_uring_free_buf_ring(&machine->ring, desc->br, desc->buf_count, bg_id);
+    rb_raise(rb_eRuntimeError, "Failed to allocate buffers");
+  }
+
+  desc->buf_mask = io_uring_buf_ring_mask(desc->buf_count);
+  void *ptr = desc->buf_base;
+  for (unsigned i = 0; i < desc->buf_count; i++) {
+		io_uring_buf_ring_add(desc->br, ptr, desc->buf_size, i, desc->buf_mask, i);
+    ptr += desc->buf_size;
+	}
+	io_uring_buf_ring_advance(desc->br, desc->buf_count);
+  machine->buffer_ring_count++;
+  return bg_id;
+}
+
+inline VALUE um_get_string_from_buffer_ring(struct um *machine, int bgid, __s32 result, __u32 flags) {
   if (!result) return Qnil;
 
   unsigned buf_idx = flags >> IORING_CQE_BUFFER_SHIFT;
@@ -64,7 +109,7 @@ inline VALUE get_string_from_buffer_ring(struct um *machine, int bgid, int resul
 
   // add buffer back to buffer ring
   io_uring_buf_ring_add(
-    desc->br, src, desc->buf_size, buf_idx, io_uring_buf_ring_mask(desc->buf_count), 0
+    desc->br, src, desc->buf_size, buf_idx, desc->buf_mask, 0
   );
   io_uring_buf_ring_advance(desc->br, 1);