Allow IO#close to interrupt fiber scheduler.

ioquatix · ioquatix · commit 5694eb120940 · 2025-05-13T22:42:26.000+09:00
diff --git a/include/ruby/fiber/scheduler.h b/include/ruby/fiber/scheduler.h
@@ -411,6 +411,12 @@ struct rb_fiber_scheduler_blocking_operation_state {
  */
 VALUE rb_fiber_scheduler_blocking_operation_wait(VALUE scheduler, void* (*function)(void *), void *data, rb_unblock_function_t *unblock_function, void *data2, int flags, struct rb_fiber_scheduler_blocking_operation_state *state);
 
+/**
+ * Interrupt a fiber by raising an exception. You can construct an exception using `rb_make_exception`.
+ *
+ */
+VALUE rb_fiber_scheduler_fiber_interrupt(VALUE scheduler, VALUE fiber, VALUE exception);
+
 /**
  * Create and schedule a non-blocking fiber.
  *
diff --git a/internal/thread.h b/internal/thread.h
@@ -72,6 +72,9 @@ void *rb_thread_prevent_fork(void *(*func)(void *), void *data); /* for ext/sock
 VALUE rb_thread_io_blocking_region(struct rb_io *io, rb_blocking_function_t *func, void *data1);
 VALUE rb_thread_io_blocking_call(struct rb_io *io, rb_blocking_function_t *func, void *data1, int events);
 
+// Invoke the given function, with the specified argument, in a way that `IO#close` from another execution context can interrupt it.
+VALUE rb_thread_io_blocking_operation(VALUE self, VALUE(*function)(VALUE), VALUE argument);
+
 /* thread.c (export) */
 int ruby_thread_has_gvl_p(void); /* for ext/fiddle/closure.c */
 
diff --git a/io_buffer.c b/io_buffer.c
@@ -2733,7 +2733,6 @@ io_buffer_blocking_region_ensure(VALUE _argument)
 static VALUE
 io_buffer_blocking_region(VALUE io, struct rb_io_buffer *buffer, rb_blocking_function_t *function, void *data)
 {
-    io = rb_io_get_io(io);
     struct rb_io *ioptr;
     RB_IO_POINTER(io, ioptr);
 
@@ -2798,6 +2797,8 @@ io_buffer_read_internal(void *_argument)
 VALUE
 rb_io_buffer_read(VALUE self, VALUE io, size_t length, size_t offset)
 {
+    io = rb_io_get_io(io);
+
     VALUE scheduler = rb_fiber_scheduler_current();
     if (scheduler != Qnil) {
         VALUE result = rb_fiber_scheduler_io_read(scheduler, io, self, length, offset);
@@ -2915,6 +2916,8 @@ io_buffer_pread_internal(void *_argument)
 VALUE
 rb_io_buffer_pread(VALUE self, VALUE io, rb_off_t from, size_t length, size_t offset)
 {
+    io = rb_io_get_io(io);
+
     VALUE scheduler = rb_fiber_scheduler_current();
     if (scheduler != Qnil) {
         VALUE result = rb_fiber_scheduler_io_pread(scheduler, io, from, self, length, offset);
@@ -3035,6 +3038,8 @@ io_buffer_write_internal(void *_argument)
 VALUE
 rb_io_buffer_write(VALUE self, VALUE io, size_t length, size_t offset)
 {
+    io = rb_io_get_write_io(rb_io_get_io(io));
+
     VALUE scheduler = rb_fiber_scheduler_current();
     if (scheduler != Qnil) {
         VALUE result = rb_fiber_scheduler_io_write(scheduler, io, self, length, offset);
@@ -3099,6 +3104,7 @@ io_buffer_write(int argc, VALUE *argv, VALUE self)
 
     return rb_io_buffer_write(self, io, length, offset);
 }
+
 struct io_buffer_pwrite_internal_argument {
     // The file descriptor to write to:
     int descriptor;
@@ -3144,6 +3150,8 @@ io_buffer_pwrite_internal(void *_argument)
 VALUE
 rb_io_buffer_pwrite(VALUE self, VALUE io, rb_off_t from, size_t length, size_t offset)
 {
+    io = rb_io_get_write_io(rb_io_get_io(io));
+
     VALUE scheduler = rb_fiber_scheduler_current();
     if (scheduler != Qnil) {
         VALUE result = rb_fiber_scheduler_io_pwrite(scheduler, io, from, self, length, offset);
diff --git a/scheduler.c b/scheduler.c
@@ -37,6 +37,7 @@ static ID id_io_close;
 static ID id_address_resolve;
 
 static ID id_blocking_operation_wait;
+static ID id_fiber_interrupt;
 
 static ID id_fiber_schedule;
 
@@ -116,6 +117,7 @@ Init_Fiber_Scheduler(void)
     id_address_resolve = rb_intern_const("address_resolve");
 
     id_blocking_operation_wait = rb_intern_const("blocking_operation_wait");
+    id_fiber_interrupt = rb_intern_const("fiber_interrupt");
 
     id_fiber_schedule = rb_intern_const("fiber");
 
@@ -442,10 +444,21 @@ rb_fiber_scheduler_unblock(VALUE scheduler, VALUE blocker, VALUE fiber)
  *  Expected to return the subset of events that are ready immediately.
  *
  */
+static VALUE
+fiber_scheduler_io_wait(VALUE _argument) {
+    VALUE *arguments = (VALUE*)_argument;
+
+    return rb_funcallv(arguments[0], id_io_wait, 3, arguments + 1);
+}
+
 VALUE
 rb_fiber_scheduler_io_wait(VALUE scheduler, VALUE io, VALUE events, VALUE timeout)
 {
-    return rb_funcall(scheduler, id_io_wait, 3, io, events, timeout);
+    VALUE arguments[] = {
+        scheduler, io, events, timeout
+    };
+
+    return rb_thread_io_blocking_operation(io, fiber_scheduler_io_wait, (VALUE)&arguments);
 }
 
 VALUE
@@ -515,14 +528,25 @@ VALUE rb_fiber_scheduler_io_selectv(VALUE scheduler, int argc, VALUE *argv)
  *
  *  The method should be considered _experimental_.
  */
+static VALUE
+fiber_scheduler_io_read(VALUE _argument) {
+    VALUE *arguments = (VALUE*)_argument;
+
+    return rb_funcallv(arguments[0], id_io_read, 4, arguments + 1);
+}
+
 VALUE
 rb_fiber_scheduler_io_read(VALUE scheduler, VALUE io, VALUE buffer, size_t length, size_t offset)
 {
+    if (!rb_respond_to(scheduler, id_io_read)) {
+        return RUBY_Qundef;
+    }
+
     VALUE arguments[] = {
-        io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
+        scheduler, io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
     };
 
-    return rb_check_funcall(scheduler, id_io_read, 4, arguments);
+    return rb_thread_io_blocking_operation(io, fiber_scheduler_io_read, (VALUE)&arguments);
 }
 
 /*
@@ -539,14 +563,25 @@ rb_fiber_scheduler_io_read(VALUE scheduler, VALUE io, VALUE buffer, size_t lengt
  *
  *  The method should be considered _experimental_.
  */
+static VALUE
+fiber_scheduler_io_pread(VALUE _argument) {
+    VALUE *arguments = (VALUE*)_argument;
+
+    return rb_funcallv(arguments[0], id_io_pread, 5, arguments + 1);
+}
+
 VALUE
 rb_fiber_scheduler_io_pread(VALUE scheduler, VALUE io, rb_off_t from, VALUE buffer, size_t length, size_t offset)
 {
+    if (!rb_respond_to(scheduler, id_io_pread)) {
+        return RUBY_Qundef;
+    }
+
     VALUE arguments[] = {
-        io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
+        scheduler, io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
     };
 
-    return rb_check_funcall(scheduler, id_io_pread, 5, arguments);
+    return rb_thread_io_blocking_operation(io, fiber_scheduler_io_pread, (VALUE)&arguments);
 }
 
 /*
@@ -577,14 +612,25 @@ rb_fiber_scheduler_io_pread(VALUE scheduler, VALUE io, rb_off_t from, VALUE buff
  *
  *  The method should be considered _experimental_.
  */
+static VALUE
+fiber_scheduler_io_write(VALUE _argument) {
+    VALUE *arguments = (VALUE*)_argument;
+
+    return rb_funcallv(arguments[0], id_io_write, 4, arguments + 1);
+}
+
 VALUE
 rb_fiber_scheduler_io_write(VALUE scheduler, VALUE io, VALUE buffer, size_t length, size_t offset)
 {
+    if (!rb_respond_to(scheduler, id_io_write)) {
+        return RUBY_Qundef;
+    }
+
     VALUE arguments[] = {
-        io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
+        scheduler, io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
     };
 
-    return rb_check_funcall(scheduler, id_io_write, 4, arguments);
+    return rb_thread_io_blocking_operation(io, fiber_scheduler_io_write, (VALUE)&arguments);
 }
 
 /*
@@ -602,14 +648,25 @@ rb_fiber_scheduler_io_write(VALUE scheduler, VALUE io, VALUE buffer, size_t leng
  *  The method should be considered _experimental_.
  *
  */
+static VALUE
+fiber_scheduler_io_pwrite(VALUE _argument) {
+    VALUE *arguments = (VALUE*)_argument;
+
+    return rb_funcallv(arguments[0], id_io_pwrite, 5, arguments + 1);
+}
+
 VALUE
 rb_fiber_scheduler_io_pwrite(VALUE scheduler, VALUE io, rb_off_t from, VALUE buffer, size_t length, size_t offset)
 {
+    if (!rb_respond_to(scheduler, id_io_pwrite)) {
+        return RUBY_Qundef;
+    }
+
     VALUE arguments[] = {
-        io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
+        scheduler, io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
     };
 
-    return rb_check_funcall(scheduler, id_io_pwrite, 5, arguments);
+    return rb_thread_io_blocking_operation(io, fiber_scheduler_io_pwrite, (VALUE)&arguments);
 }
 
 VALUE
@@ -766,6 +823,15 @@ VALUE rb_fiber_scheduler_blocking_operation_wait(VALUE scheduler, void* (*functi
     return rb_check_funcall(scheduler, id_blocking_operation_wait, 1, &proc);
 }
 
+VALUE rb_fiber_scheduler_fiber_interrupt(VALUE scheduler, VALUE fiber, VALUE exception)
+{
+    VALUE arguments[] = {
+        fiber, exception
+    };
+
+    return rb_check_funcall(scheduler, id_fiber_interrupt, 2, arguments);
+}
+
 /*
  *  Document-method: Fiber::Scheduler#fiber
  *  call-seq: fiber(&block)
diff --git a/test/fiber/scheduler.rb b/test/fiber/scheduler.rb
@@ -68,9 +68,15 @@ def next_timeout
   def run
     # $stderr.puts [__method__, Fiber.current].inspect
 
+    readable = writable = nil
+
     while @readable.any? or @writable.any? or @waiting.any? or @blocking.any?
       # May only handle file descriptors up to 1024...
-      readable, writable = IO.select(@readable.keys + [@urgent.first], @writable.keys, [], next_timeout)
+      begin
+        readable, writable = IO.select(@readable.keys + [@urgent.first], @writable.keys, [], next_timeout)
+      rescue IOError
+        # Ignore - this can happen if the IO is closed while we are waiting.
+      end
 
       # puts "readable: #{readable}" if readable&.any?
       # puts "writable: #{writable}" if writable&.any?
@@ -290,6 +296,30 @@ def unblock(blocker, fiber)
     io.write_nonblock('.')
   end
 
+  class FiberInterrupt
+    def initialize(fiber, exception)
+      @fiber = fiber
+      @exception = exception
+    end
+
+    def alive?
+      @fiber.alive?
+    end
+
+    def transfer
+      @fiber.raise(@exception)
+    end
+  end
+
+  def fiber_interrupt(fiber, exception)
+    @lock.synchronize do
+      @ready << FiberInterrupt.new(fiber, exception)
+    end
+
+    io = @urgent.last
+    io.write_nonblock('.')
+  end
+
   # This hook is invoked by `Fiber.schedule`. Strictly speaking, you should use
   # it to create scheduled fibers, but it is not required in practice;
   # `Fiber.new` is usually sufficient.
diff --git a/thread.c b/thread.c
@@ -1731,6 +1731,37 @@ rb_io_blocking_operation_release(struct rb_io *io, struct rb_io_blocking_operati
     }
 }
 
+
+static VALUE
+rb_thread_io_blocking_operation_ensure(VALUE _argument)
+{
+    struct io_blocking_operation_arguments *arguments = (void*)_argument;
+
+    rb_io_blocking_operation_release(arguments->io, arguments->blocking_operation);
+
+    return Qnil;
+}
+
+VALUE
+rb_thread_io_blocking_operation(VALUE self, VALUE(*function)(VALUE), VALUE argument)
+{
+    struct rb_io *io;
+    RB_IO_POINTER(self, io);
+
+    rb_execution_context_t *ec = GET_EC();
+    struct rb_io_blocking_operation blocking_operation = {
+        .ec = ec,
+    };
+    ccan_list_add(&io->blocking_operations, &blocking_operation.list);
+
+    struct io_blocking_operation_arguments io_blocking_operation_arguments = {
+        .io = io,
+        .blocking_operation = &blocking_operation
+    };
+
+    return rb_ensure(function, argument, rb_thread_io_blocking_operation_ensure, (VALUE)&io_blocking_operation_arguments);
+}
+
 static bool
 thread_io_mn_schedulable(rb_thread_t *th, int events, const struct timeval *timeout)
 {
@@ -1832,7 +1863,7 @@ rb_thread_io_blocking_call(struct rb_io* io, rb_blocking_function_t *func, void
                 saved_errno = errno;
             }, ubf_select, th, FALSE);
 
-            th = rb_ec_thread_ptr(ec);
+            RUBY_ASSERT(th == rb_ec_thread_ptr(ec));
             if (events &&
                 blocking_call_retryable_p((int)val, saved_errno) &&
                 thread_io_wait_events(th, fd, events, NULL)) {
@@ -2636,10 +2667,10 @@ rb_ec_reset_raised(rb_execution_context_t *ec)
     return 1;
 }
 
-static size_t
-thread_io_close_notify_all(struct rb_io *io)
+static VALUE
+thread_io_close_notify_all(VALUE _io)
 {
-    RUBY_ASSERT_CRITICAL_SECTION_ENTER();
+    struct rb_io *io = (struct rb_io *)_io;
 
     size_t count = 0;
     rb_vm_t *vm = io->closing_ec->thread_ptr->vm;
@@ -2650,17 +2681,17 @@ thread_io_close_notify_all(struct rb_io *io)
         rb_execution_context_t *ec = blocking_operation->ec;
 
         rb_thread_t *thread = ec->thread_ptr;
-        rb_threadptr_pending_interrupt_enque(thread, error);
-
-        // This operation is slow:
-        rb_threadptr_interrupt(thread);
+        if (thread->scheduler != Qnil) {
+            rb_fiber_scheduler_fiber_interrupt(thread->scheduler, rb_fiberptr_self(ec->fiber_ptr), error);
+        } else {
+            rb_threadptr_pending_interrupt_enque(thread, error);
+            rb_threadptr_interrupt(thread);
+        }
 
         count += 1;
     }
 
-    RUBY_ASSERT_CRITICAL_SECTION_LEAVE();
-
-    return count;
+    return (VALUE)count;
 }
 
 size_t
@@ -2683,7 +2714,9 @@ rb_thread_io_close_interrupt(struct rb_io *io)
     // This is used to ensure the correct execution context is woken up after the blocking operation is interrupted:
     io->wakeup_mutex = rb_mutex_new();
 
-    return thread_io_close_notify_all(io);
+    VALUE result = rb_mutex_synchronize(io->wakeup_mutex, thread_io_close_notify_all, (VALUE)io);
+
+    return (size_t)result;
 }
 
 void
