polyphony 0.55.0 → 0.59.1

data/ext/polyphony/backend_io_uring.c

@@ -4,7 +4,6 @@
 #include <sys/socket.h>
 #include <sys/uio.h>
 #include <unistd.h>
-#include <fcntl.h>
 #include <netinet/in.h>
 #include <arpa/inet.h>
 #include <stdnoreturn.h>
@@ -19,6 +18,7 @@
 #include "backend_io_uring_context.h"
 #include "ruby/thread.h"
 #include "ruby/io.h"
+#include "backend_common.h"
 
 VALUE SYM_io_uring;
 VALUE SYM_send;
@@ -26,34 +26,13 @@ VALUE SYM_splice;
 VALUE SYM_write;
 
 #ifdef POLYPHONY_UNSET_NONBLOCK
-ID ID_ivar_is_nonblocking;
-
-// One of the changes introduced in Ruby 3.0 as part of the work on the
-// FiberScheduler interface is that all created sockets are marked as
-// non-blocking. This prevents the io_uring backend from working correctly,
-// since it will return an EAGAIN error just like a normal syscall. So here
-// instead of setting O_NONBLOCK (which is required for the libev backend), we
-// unset it.
-inline void io_unset_nonblock(rb_io_t *fptr, VALUE io) {
-  VALUE is_nonblocking = rb_ivar_get(io, ID_ivar_is_nonblocking);
-  if (is_nonblocking == Qfalse) return;
-
-  rb_ivar_set(io, ID_ivar_is_nonblocking, Qfalse);
-
-  int oflags = fcntl(fptr->fd, F_GETFL);
-  if ((oflags == -1) && (oflags & O_NONBLOCK)) return;
-  oflags &= !O_NONBLOCK;
-  fcntl(fptr->fd, F_SETFL, oflags);
-}
+#define io_unset_nonblock(fptr, io) io_verify_blocking_mode(fptr, io, Qtrue)
 #else
-// NOP
 #define io_unset_nonblock(fptr, io)
 #endif
 
 typedef struct Backend_t {
-  // common fields
-  unsigned int        currently_polling;
-  unsigned int        pending_count;
+  struct Backend_base base;
 
   // implementation-specific fields
   struct io_uring     ring;
@@ -63,7 +42,15 @@ typedef struct Backend_t {
   int                 event_fd;
 } Backend_t;
 
-#include "backend_common.h"
+static void Backend_mark(void *ptr) {
+  Backend_t *backend = ptr;
+  backend_base_mark(&backend->base);
+}
+
+static void Backend_free(void *ptr) {
+  Backend_t *backend = ptr;
+  backend_base_finalize(&backend->base);
+}
 
 static size_t Backend_size(const void *ptr) {
   return sizeof(Backend_t);
@@ -71,7 +58,7 @@ static size_t Backend_size(const void *ptr) {
 
 static const rb_data_type_t Backend_type = {
     "IOUringBackend",
-    {0, 0, Backend_size,},
+    {Backend_mark, Backend_free, Backend_size,},
     0, 0, RUBY_TYPED_FREE_IMMEDIATELY
 };
 
@@ -88,8 +75,7 @@ static VALUE Backend_initialize(VALUE self) {
   Backend_t *backend;
   GetBackend(self, backend);
 
-  backend->currently_polling = 0;
-  backend->pending_count = 0;
+  backend_base_initialize(&backend->base);
   backend->pending_sqes = 0;
   backend->prepared_limit = 2048;
 
@@ -117,20 +103,13 @@ VALUE Backend_post_fork(VALUE self) {
   io_uring_queue_exit(&backend->ring);
   io_uring_queue_init(backend->prepared_limit, &backend->ring, 0);
   context_store_free(&backend->store);
-  backend->currently_polling = 0;
-  backend->pending_count = 0;
+  backend->base.currently_polling = 0;
+  backend->base.pending_count = 0;
   backend->pending_sqes = 0;
 
   return self;
 }
 
-unsigned int Backend_pending_count(VALUE self) {
-  Backend_t *backend;
-  GetBackend(self, backend);
-
-  return backend->pending_count;
-}
-
 typedef struct poll_context {
   struct io_uring     *ring;
   struct io_uring_cqe *cqe;
@@ -154,6 +133,7 @@ static inline void io_uring_backend_handle_completion(struct io_uring_cqe *cqe,
   op_context_t *ctx = io_uring_cqe_get_data(cqe);
   if (!ctx) return;
 
+  // printf("cqe ctx %p id: %d result: %d (%s, ref_count: %d)\n", ctx, ctx->id, cqe->res, op_type_to_str(ctx->type), ctx->ref_count);
   ctx->result = cqe->res;
   if (ctx->ref_count == 2 && ctx->result != -ECANCELED && ctx->fiber)
     Fiber_make_runnable(ctx->fiber, ctx->resume_value);
@@ -197,38 +177,70 @@ void io_uring_backend_poll(Backend_t *backend) {
     io_uring_submit(&backend->ring);
   }
 
-  backend->currently_polling = 1;
+  backend->base.currently_polling = 1;
   rb_thread_call_without_gvl(io_uring_backend_poll_without_gvl, (void *)&poll_ctx, RUBY_UBF_IO, 0);
-  backend->currently_polling = 0;
+  backend->base.currently_polling = 0;
   if (poll_ctx.result < 0) return;
 
   io_uring_backend_handle_completion(poll_ctx.cqe, backend);
   io_uring_cqe_seen(&backend->ring, poll_ctx.cqe);
 }
 
-VALUE Backend_poll(VALUE self, VALUE nowait, VALUE current_fiber, VALUE runqueue) {
-  int is_nowait = nowait == Qtrue;
+inline VALUE Backend_poll(VALUE self, VALUE blocking) {
+  int is_blocking = blocking == Qtrue;
   Backend_t *backend;
   GetBackend(self, backend);
 
-  if (is_nowait && backend->pending_sqes) {
+  if (!is_blocking && backend->pending_sqes) {
     backend->pending_sqes = 0;
     io_uring_submit(&backend->ring);
   }
 
-  COND_TRACE(2, SYM_fiber_event_poll_enter, current_fiber);
-  if (!is_nowait) io_uring_backend_poll(backend);
+  COND_TRACE(&backend->base, 2, SYM_fiber_event_poll_enter, rb_fiber_current());
+  if (is_blocking) io_uring_backend_poll(backend);
   io_uring_backend_handle_ready_cqes(backend);
-  COND_TRACE(2, SYM_fiber_event_poll_leave, current_fiber);
+  COND_TRACE(&backend->base, 2, SYM_fiber_event_poll_leave, rb_fiber_current());
   
   return self;
 }
 
+inline void Backend_schedule_fiber(VALUE thread, VALUE self, VALUE fiber, VALUE value, int prioritize) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+
+  backend_base_schedule_fiber(thread, self, &backend->base, fiber, value, prioritize);
+}
+
+inline void Backend_unschedule_fiber(VALUE self, VALUE fiber) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+
+  runqueue_delete(&backend->base.runqueue, fiber);  
+}
+
+inline VALUE Backend_switch_fiber(VALUE self) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+
+  return backend_base_switch_fiber(self, &backend->base);
+}
+
+inline struct backend_stats Backend_stats(VALUE self) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+
+  return (struct backend_stats){
+    .scheduled_fibers = runqueue_len(&backend->base.runqueue),
+    .waiting_fibers = 0,
+    .pending_ops = backend->base.pending_count
+  };
+}
+
 VALUE Backend_wakeup(VALUE self) {
   Backend_t *backend;
   GetBackend(self, backend);
 
-  if (backend->currently_polling) {
+  if (backend->base.currently_polling) {
     // Since we're currently blocking while waiting for a completion, we add a
     // NOP which would cause the io_uring_enter syscall to return
     struct io_uring_sqe *sqe = io_uring_get_sqe(&backend->ring);
@@ -259,11 +271,13 @@ int io_uring_backend_defer_submit_and_await(
 {
   VALUE switchpoint_result = Qnil;
 
-  io_uring_sqe_set_data(sqe, ctx);
-  io_uring_sqe_set_flags(sqe, IOSQE_ASYNC);
+  if (sqe) {
+    io_uring_sqe_set_data(sqe, ctx);
+    io_uring_sqe_set_flags(sqe, IOSQE_ASYNC);
+  }
   io_uring_backend_defer_submit(backend);
 
-  switchpoint_result = backend_await(backend);
+  switchpoint_result = backend_await((struct Backend_base *)backend);
 
   if (ctx->ref_count > 1) {
     // op was not completed (an exception was raised), so we need to cancel it
@@ -466,7 +480,7 @@ VALUE Backend_write(VALUE self, VALUE io, VALUE str) {
     VALUE resume_value = Qnil;
     op_context_t *ctx = context_store_acquire(&backend->store, OP_WRITE);
     struct io_uring_sqe *sqe = io_uring_get_sqe(&backend->ring);
-    io_uring_prep_write(sqe, fptr->fd, buf, left, -1);
+    io_uring_prep_write(sqe, fptr->fd, buf, left, 0);
     
     int result = io_uring_backend_defer_submit_and_await(backend, sqe, ctx, &resume_value);
     int completed = context_store_release(&backend->store, ctx);
@@ -926,7 +940,6 @@ int io_uring_backend_submit_timeout_and_await(Backend_t *backend, double duratio
   
   op_context_t *ctx = context_store_acquire(&backend->store, OP_TIMEOUT);
   io_uring_prep_timeout(sqe, &ts, 0, 0);
-
   io_uring_backend_defer_submit_and_await(backend, sqe, ctx, resume_value);
   return context_store_release(&backend->store, ctx);
 }
@@ -1085,7 +1098,7 @@ struct io_uring_sqe *Backend_chain_prepare_write(Backend_t *backend, VALUE io, V
   long len = RSTRING_LEN(str);
 
   struct io_uring_sqe *sqe = io_uring_get_sqe(&backend->ring);
-  io_uring_prep_write(sqe, fptr->fd, buf, len, -1);
+  io_uring_prep_write(sqe, fptr->fd, buf, len, 0);
   return sqe;
 }
 
@@ -1177,7 +1190,7 @@ VALUE Backend_chain(int argc,VALUE *argv, VALUE self) {
 
   ctx->ref_count = sqe_count + 1;
   io_uring_backend_defer_submit(backend);
-  resume_value = backend_await(backend);
+  resume_value = backend_await((struct Backend_base *)backend);
   int result = ctx->result;
   int completed = context_store_release(&backend->store, ctx);
   if (!completed) {
@@ -1195,6 +1208,197 @@ VALUE Backend_chain(int argc,VALUE *argv, VALUE self) {
   return INT2NUM(result);
 }
 
+VALUE Backend_idle_gc_period_set(VALUE self, VALUE period) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+  backend->base.idle_gc_period = NUM2DBL(period);
+  backend->base.idle_gc_last_time = current_time();
+  return self;
+}
+
+VALUE Backend_idle_proc_set(VALUE self, VALUE block) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+  backend->base.idle_proc = block;
+  return self;
+}
+
+inline VALUE Backend_run_idle_tasks(VALUE self) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+  backend_run_idle_tasks(&backend->base);
+  return self;
+}
+
+static inline void splice_chunks_prep_write(op_context_t *ctx, struct io_uring_sqe *sqe, int fd, VALUE str) {
+  char *buf = RSTRING_PTR(str);
+  int len = RSTRING_LEN(str);
+  io_uring_prep_write(sqe, fd, buf, len, 0);
+  // io_uring_prep_send(sqe, fd, buf, len, 0);
+  io_uring_sqe_set_data(sqe, ctx);
+}
+
+static inline void splice_chunks_prep_splice(op_context_t *ctx, struct io_uring_sqe *sqe, int src, int dest, int maxlen) {
+  io_uring_prep_splice(sqe, src, -1, dest, -1, maxlen, 0);
+  io_uring_sqe_set_data(sqe, ctx);
+}
+
+static inline void splice_chunks_get_sqe(
+  Backend_t *backend,
+  op_context_t **ctx,
+  struct io_uring_sqe **sqe,
+  enum op_type type
+)
+{
+  if (*ctx) {
+    if (*sqe) (*sqe)->flags |= IOSQE_IO_LINK;
+    (*ctx)->ref_count++;
+  }
+  else
+    *ctx = context_store_acquire(&backend->store, type);
+  (*sqe) = io_uring_get_sqe(&backend->ring);
+}
+
+static inline void splice_chunks_cancel(Backend_t *backend, op_context_t *ctx) {
+  ctx->result = -ECANCELED;
+  struct io_uring_sqe *sqe = io_uring_get_sqe(&backend->ring);
+  io_uring_prep_cancel(sqe, ctx, 0);
+  backend->pending_sqes = 0;
+  io_uring_submit(&backend->ring);
+}
+
+static inline int splice_chunks_await_ops(
+  Backend_t *backend,
+  op_context_t **ctx,
+  int *result,
+  VALUE *switchpoint_result
+)
+{
+  int res = io_uring_backend_defer_submit_and_await(backend, 0, *ctx, switchpoint_result);
+  if (result) (*result) = res;
+  int completed = context_store_release(&backend->store, *ctx);
+  if (!completed) {
+    splice_chunks_cancel(backend, *ctx);
+    if (TEST_EXCEPTION(*switchpoint_result)) return 1;
+  }
+  *ctx = 0;
+  return 0;
+}
+
+#define SPLICE_CHUNKS_AWAIT_OPS(backend, ctx, result, switchpoint_result) \
+  if (splice_chunks_await_ops(backend, ctx, result, switchpoint_result)) goto error;
+
+VALUE Backend_splice_chunks(VALUE self, VALUE src, VALUE dest, VALUE prefix, VALUE postfix, VALUE chunk_prefix, VALUE chunk_postfix, VALUE chunk_size) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+  int total = 0;
+  int err = 0;
+  VALUE switchpoint_result = Qnil;
+  op_context_t *ctx = 0;
+  struct io_uring_sqe *sqe = 0;
+
+  rb_io_t *src_fptr;
+  rb_io_t *dest_fptr;
+
+  VALUE underlying_io = rb_ivar_get(src, ID_ivar_io);
+  if (underlying_io != Qnil) src = underlying_io;
+  GetOpenFile(src, src_fptr);
+  io_verify_blocking_mode(src_fptr, src, Qtrue);
+
+  underlying_io = rb_ivar_get(dest, ID_ivar_io);
+  if (underlying_io != Qnil) dest = underlying_io;
+  dest = rb_io_get_write_io(dest);
+  GetOpenFile(dest, dest_fptr);
+  io_verify_blocking_mode(dest_fptr, dest, Qtrue);
+
+  int maxlen = NUM2INT(chunk_size);
+  VALUE str = Qnil;
+  VALUE chunk_len_value = Qnil;
+
+  int pipefd[2] = { -1, -1 };
+  if (pipe(pipefd) == -1) {
+    err = errno;
+    goto syscallerror;
+  }
+
+  if (prefix != Qnil) {
+    splice_chunks_get_sqe(backend, &ctx, &sqe, OP_WRITE);
+    splice_chunks_prep_write(ctx, sqe, dest_fptr->fd, prefix);
+  }
+
+  while (1) {
+    int chunk_len;
+    VALUE chunk_prefix_str = Qnil;
+    VALUE chunk_postfix_str = Qnil;
+
+    splice_chunks_get_sqe(backend, &ctx, &sqe, OP_SPLICE);
+    splice_chunks_prep_splice(ctx, sqe, src_fptr->fd, pipefd[1], maxlen);
+    
+    SPLICE_CHUNKS_AWAIT_OPS(backend, &ctx, &chunk_len, &switchpoint_result);
+    if (chunk_len == 0) break;
+    
+    total += chunk_len;
+    chunk_len_value = INT2NUM(chunk_len);
+
+
+    if (chunk_prefix != Qnil) {
+      chunk_prefix_str = (TYPE(chunk_prefix) == T_STRING) ? chunk_prefix : rb_funcall(chunk_prefix, ID_call, 1, chunk_len_value);
+      splice_chunks_get_sqe(backend, &ctx, &sqe, OP_WRITE);
+      splice_chunks_prep_write(ctx, sqe, dest_fptr->fd, chunk_prefix_str);
+    }
+
+    splice_chunks_get_sqe(backend, &ctx, &sqe, OP_SPLICE);
+    splice_chunks_prep_splice(ctx, sqe, pipefd[0], dest_fptr->fd, chunk_len);
+
+    if (chunk_postfix != Qnil) {
+      chunk_postfix_str = (TYPE(chunk_postfix) == T_STRING) ? chunk_postfix : rb_funcall(chunk_postfix, ID_call, 1, chunk_len_value);
+      splice_chunks_get_sqe(backend, &ctx, &sqe, OP_WRITE);
+      splice_chunks_prep_write(ctx, sqe, dest_fptr->fd, chunk_postfix_str);
+    }
+
+    RB_GC_GUARD(chunk_prefix_str);
+    RB_GC_GUARD(chunk_postfix_str);
+  }
+
+  if (postfix != Qnil) {
+    splice_chunks_get_sqe(backend, &ctx, &sqe, OP_WRITE);
+    splice_chunks_prep_write(ctx, sqe, dest_fptr->fd, postfix);
+  }
+  if (ctx) {
+    SPLICE_CHUNKS_AWAIT_OPS(backend, &ctx, 0, &switchpoint_result);
+  }
+
+  RB_GC_GUARD(str);
+  RB_GC_GUARD(chunk_len_value);
+  RB_GC_GUARD(switchpoint_result);
+  if (pipefd[0] != -1) close(pipefd[0]);
+  if (pipefd[1] != -1) close(pipefd[1]);
+  return INT2NUM(total);
+syscallerror:
+  if (pipefd[0] != -1) close(pipefd[0]);
+  if (pipefd[1] != -1) close(pipefd[1]);
+  rb_syserr_fail(err, strerror(err));
+error:
+  if (pipefd[0] != -1) close(pipefd[0]);
+  if (pipefd[1] != -1) close(pipefd[1]);
+  return RAISE_EXCEPTION(switchpoint_result);
+}
+
+VALUE Backend_trace(int argc, VALUE *argv, VALUE self) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+  backend_trace(&backend->base, argc, argv);
+  return self;
+}
+
+VALUE Backend_trace_proc_set(VALUE self, VALUE block) {
+  Backend_t *backend;
+  GetBackend(self, backend);
+
+  backend->base.trace_proc = block;
+  return self;
+}
+
 void Init_Backend() {
   VALUE cBackend = rb_define_class_under(mPolyphony, "Backend", rb_cObject);
   rb_define_alloc_func(cBackend, Backend_allocate);
@@ -1202,11 +1406,16 @@ void Init_Backend() {
   rb_define_method(cBackend, "initialize", Backend_initialize, 0);
   rb_define_method(cBackend, "finalize", Backend_finalize, 0);
   rb_define_method(cBackend, "post_fork", Backend_post_fork, 0);
+  rb_define_method(cBackend, "trace", Backend_trace, -1);
+  rb_define_method(cBackend, "trace_proc=", Backend_trace_proc_set, 1);
 
-  rb_define_method(cBackend, "poll", Backend_poll, 3);
+  rb_define_method(cBackend, "poll", Backend_poll, 1);
   rb_define_method(cBackend, "break", Backend_wakeup, 0);
   rb_define_method(cBackend, "kind", Backend_kind, 0);
   rb_define_method(cBackend, "chain", Backend_chain, -1);
+  rb_define_method(cBackend, "idle_gc_period=", Backend_idle_gc_period_set, 1);
+  rb_define_method(cBackend, "idle_proc=", Backend_idle_proc_set, 1);
+  rb_define_method(cBackend, "splice_chunks", Backend_splice_chunks, 7);
 
   rb_define_method(cBackend, "accept", Backend_accept, 2);
   rb_define_method(cBackend, "accept_loop", Backend_accept_loop, 2);