Skip to content

Commit 56d24c9

Browse files
ioquatixioquatix
committed
Allow IO#close to interrupt fiber scheduler.
1 parent 12e2dfd commit 56d24c9

File tree

6 files changed

+175
-31
lines changed

6 files changed

+175
-31
lines changed

‎include/ruby/fiber/scheduler.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,12 @@ struct rb_fiber_scheduler_blocking_operation_state {
411411
*/
412412
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);
413413

414+
/**
415+
* Interrupt a fiber by raising an exception. You can construct an exception using `rb_make_exception`.
416+
*
417+
*/
418+
VALUE rb_fiber_scheduler_fiber_interrupt(VALUE scheduler, VALUE fiber, VALUE exception);
419+
414420
/**
415421
* Create and schedule a non-blocking fiber.
416422
*

‎internal/thread.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,9 @@ void *rb_thread_prevent_fork(void *(*func)(void *), void *data); /* for ext/sock
7676
VALUE rb_thread_io_blocking_region(rb_blocking_function_t *func, void *data1, int fd);
7777
VALUE rb_thread_io_blocking_call(rb_blocking_function_t *func, void *data1, int fd, int events);
7878

79+
// Invoke the given function, with the specified argument, in a way that `IO#close` from another execution context can interrupt it.
80+
VALUE rb_thread_io_interruptible_operation(VALUE self, VALUE(*function)(VALUE), VALUE argument);
81+
7982
/* thread.c (export) */
8083
int ruby_thread_has_gvl_p(void); /* for ext/fiddle/closure.c */
8184

‎io_buffer.c‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2735,7 +2735,6 @@ io_buffer_blocking_region_ensure(VALUE _argument)
27352735
static VALUE
27362736
io_buffer_blocking_region(VALUE io, struct rb_io_buffer *buffer, rb_blocking_function_t *function, void *data)
27372737
{
2738-
io = rb_io_get_io(io);
27392738
struct rb_io *ioptr;
27402739
RB_IO_POINTER(io, ioptr);
27412740

@@ -2800,6 +2799,8 @@ io_buffer_read_internal(void *_argument)
28002799
VALUE
28012800
rb_io_buffer_read(VALUE self, VALUE io, size_t length, size_t offset)
28022801
{
2802+
io = rb_io_get_io(io);
2803+
28032804
VALUE scheduler = rb_fiber_scheduler_current();
28042805
if (scheduler != Qnil) {
28052806
VALUE result = rb_fiber_scheduler_io_read(scheduler, io, self, length, offset);
@@ -2917,6 +2918,8 @@ io_buffer_pread_internal(void *_argument)
29172918
VALUE
29182919
rb_io_buffer_pread(VALUE self, VALUE io, rb_off_t from, size_t length, size_t offset)
29192920
{
2921+
io = rb_io_get_io(io);
2922+
29202923
VALUE scheduler = rb_fiber_scheduler_current();
29212924
if (scheduler != Qnil) {
29222925
VALUE result = rb_fiber_scheduler_io_pread(scheduler, io, from, self, length, offset);
@@ -3037,6 +3040,8 @@ io_buffer_write_internal(void *_argument)
30373040
VALUE
30383041
rb_io_buffer_write(VALUE self, VALUE io, size_t length, size_t offset)
30393042
{
3043+
io = rb_io_get_write_io(rb_io_get_io(io));
3044+
30403045
VALUE scheduler = rb_fiber_scheduler_current();
30413046
if (scheduler != Qnil) {
30423047
VALUE result = rb_fiber_scheduler_io_write(scheduler, io, self, length, offset);
@@ -3101,6 +3106,7 @@ io_buffer_write(int argc, VALUE *argv, VALUE self)
31013106

31023107
return rb_io_buffer_write(self, io, length, offset);
31033108
}
3109+
31043110
struct io_buffer_pwrite_internal_argument {
31053111
// The file descriptor to write to:
31063112
int descriptor;
@@ -3146,6 +3152,8 @@ io_buffer_pwrite_internal(void *_argument)
31463152
VALUE
31473153
rb_io_buffer_pwrite(VALUE self, VALUE io, rb_off_t from, size_t length, size_t offset)
31483154
{
3155+
io = rb_io_get_write_io(rb_io_get_io(io));
3156+
31493157
VALUE scheduler = rb_fiber_scheduler_current();
31503158
if (scheduler != Qnil) {
31513159
VALUE result = rb_fiber_scheduler_io_pwrite(scheduler, io, from, self, length, offset);

‎scheduler.c‎

Lines changed: 75 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ static ID id_io_close;
3737
static ID id_address_resolve;
3838

3939
static ID id_blocking_operation_wait;
40+
static ID id_fiber_interrupt;
4041

4142
static ID id_fiber_schedule;
4243

@@ -116,6 +117,7 @@ Init_Fiber_Scheduler(void)
116117
id_address_resolve = rb_intern_const("address_resolve");
117118

118119
id_blocking_operation_wait = rb_intern_const("blocking_operation_wait");
120+
id_fiber_interrupt = rb_intern_const("fiber_interrupt");
119121

120122
id_fiber_schedule = rb_intern_const("fiber");
121123

@@ -442,10 +444,21 @@ rb_fiber_scheduler_unblock(VALUE scheduler, VALUE blocker, VALUE fiber)
442444
* Expected to return the subset of events that are ready immediately.
443445
*
444446
*/
447+
static VALUE
448+
fiber_scheduler_io_wait(VALUE _argument) {
449+
VALUE *arguments = (VALUE*)_argument;
450+
451+
return rb_funcallv(arguments[0], id_io_wait, 3, arguments + 1);
452+
}
453+
445454
VALUE
446455
rb_fiber_scheduler_io_wait(VALUE scheduler, VALUE io, VALUE events, VALUE timeout)
447456
{
448-
return rb_funcall(scheduler, id_io_wait, 3, io, events, timeout);
457+
VALUE arguments[] = {
458+
scheduler, io, events, timeout
459+
};
460+
461+
return rb_thread_io_interruptible_operation(io, fiber_scheduler_io_wait, (VALUE)&arguments);
449462
}
450463

451464
VALUE
@@ -515,14 +528,25 @@ VALUE rb_fiber_scheduler_io_selectv(VALUE scheduler, int argc, VALUE *argv)
515528
*
516529
* The method should be considered _experimental_.
517530
*/
531+
static VALUE
532+
fiber_scheduler_io_read(VALUE _argument) {
533+
VALUE *arguments = (VALUE*)_argument;
534+
535+
return rb_funcallv(arguments[0], id_io_read, 4, arguments + 1);
536+
}
537+
518538
VALUE
519539
rb_fiber_scheduler_io_read(VALUE scheduler, VALUE io, VALUE buffer, size_t length, size_t offset)
520540
{
541+
if (!rb_respond_to(scheduler, id_io_read)) {
542+
return RUBY_Qundef;
543+
}
544+
521545
VALUE arguments[] = {
522-
io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
546+
scheduler, io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
523547
};
524548

525-
return rb_check_funcall(scheduler, id_io_read, 4, arguments);
549+
return rb_thread_io_interruptible_operation(io, fiber_scheduler_io_read, (VALUE)&arguments);
526550
}
527551

528552
/*
@@ -539,14 +563,25 @@ rb_fiber_scheduler_io_read(VALUE scheduler, VALUE io, VALUE buffer, size_t lengt
539563
*
540564
* The method should be considered _experimental_.
541565
*/
566+
static VALUE
567+
fiber_scheduler_io_pread(VALUE _argument) {
568+
VALUE *arguments = (VALUE*)_argument;
569+
570+
return rb_funcallv(arguments[0], id_io_pread, 5, arguments + 1);
571+
}
572+
542573
VALUE
543574
rb_fiber_scheduler_io_pread(VALUE scheduler, VALUE io, rb_off_t from, VALUE buffer, size_t length, size_t offset)
544575
{
576+
if (!rb_respond_to(scheduler, id_io_pread)) {
577+
return RUBY_Qundef;
578+
}
579+
545580
VALUE arguments[] = {
546-
io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
581+
scheduler, io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
547582
};
548583

549-
return rb_check_funcall(scheduler, id_io_pread, 5, arguments);
584+
return rb_thread_io_interruptible_operation(io, fiber_scheduler_io_pread, (VALUE)&arguments);
550585
}
551586

552587
/*
@@ -577,14 +612,25 @@ rb_fiber_scheduler_io_pread(VALUE scheduler, VALUE io, rb_off_t from, VALUE buff
577612
*
578613
* The method should be considered _experimental_.
579614
*/
615+
static VALUE
616+
fiber_scheduler_io_write(VALUE _argument) {
617+
VALUE *arguments = (VALUE*)_argument;
618+
619+
return rb_funcallv(arguments[0], id_io_write, 4, arguments + 1);
620+
}
621+
580622
VALUE
581623
rb_fiber_scheduler_io_write(VALUE scheduler, VALUE io, VALUE buffer, size_t length, size_t offset)
582624
{
625+
if (!rb_respond_to(scheduler, id_io_write)) {
626+
return RUBY_Qundef;
627+
}
628+
583629
VALUE arguments[] = {
584-
io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
630+
scheduler, io, buffer, SIZET2NUM(length), SIZET2NUM(offset)
585631
};
586632

587-
return rb_check_funcall(scheduler, id_io_write, 4, arguments);
633+
return rb_thread_io_interruptible_operation(io, fiber_scheduler_io_write, (VALUE)&arguments);
588634
}
589635

590636
/*
@@ -602,14 +648,25 @@ rb_fiber_scheduler_io_write(VALUE scheduler, VALUE io, VALUE buffer, size_t leng
602648
* The method should be considered _experimental_.
603649
*
604650
*/
651+
static VALUE
652+
fiber_scheduler_io_pwrite(VALUE _argument) {
653+
VALUE *arguments = (VALUE*)_argument;
654+
655+
return rb_funcallv(arguments[0], id_io_pwrite, 5, arguments + 1);
656+
}
657+
605658
VALUE
606659
rb_fiber_scheduler_io_pwrite(VALUE scheduler, VALUE io, rb_off_t from, VALUE buffer, size_t length, size_t offset)
607660
{
661+
if (!rb_respond_to(scheduler, id_io_pwrite)) {
662+
return RUBY_Qundef;
663+
}
664+
608665
VALUE arguments[] = {
609-
io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
666+
scheduler, io, buffer, OFFT2NUM(from), SIZET2NUM(length), SIZET2NUM(offset)
610667
};
611668

612-
return rb_check_funcall(scheduler, id_io_pwrite, 5, arguments);
669+
return rb_thread_io_interruptible_operation(io, fiber_scheduler_io_pwrite, (VALUE)&arguments);
613670
}
614671

615672
VALUE
@@ -766,6 +823,15 @@ VALUE rb_fiber_scheduler_blocking_operation_wait(VALUE scheduler, void* (*functi
766823
return rb_check_funcall(scheduler, id_blocking_operation_wait, 1, &proc);
767824
}
768825

826+
VALUE rb_fiber_scheduler_fiber_interrupt(VALUE scheduler, VALUE fiber, VALUE exception)
827+
{
828+
VALUE arguments[] = {
829+
fiber, exception
830+
};
831+
832+
return rb_check_funcall(scheduler, id_fiber_interrupt, 2, arguments);
833+
}
834+
769835
/*
770836
* Document-method: Fiber::Scheduler#fiber
771837
* call-seq: fiber(&block)

‎test/fiber/scheduler.rb‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,15 @@ def next_timeout
6868
def run
6969
# $stderr.puts [__method__, Fiber.current].inspect
7070

71+
readable = writable = nil
72+
7173
while @readable.any? or @writable.any? or @waiting.any? or @blocking.any?
7274
# May only handle file descriptors up to 1024...
73-
readable, writable = IO.select(@readable.keys + [@urgent.first], @writable.keys, [], next_timeout)
75+
begin
76+
readable, writable = IO.select(@readable.keys + [@urgent.first], @writable.keys, [], next_timeout)
77+
rescue IOError
78+
# Ignore - this can happen if the IO is closed while we are waiting.
79+
end
7480

7581
# puts "readable: #{readable}" if readable&.any?
7682
# puts "writable: #{writable}" if writable&.any?
@@ -290,6 +296,30 @@ def unblock(blocker, fiber)
290296
io.write_nonblock('.')
291297
end
292298

299+
class FiberInterrupt
300+
def initialize(fiber, exception)
301+
@fiber = fiber
302+
@exception = exception
303+
end
304+
305+
def alive?
306+
@fiber.alive?
307+
end
308+
309+
def transfer
310+
@fiber.raise(@exception)
311+
end
312+
end
313+
314+
def fiber_interrupt(fiber, exception)
315+
@lock.synchronize do
316+
@ready << FiberInterrupt.new(fiber, exception)
317+
end
318+
319+
io = @urgent.last
320+
io.write_nonblock('.')
321+
end
322+
293323
# This hook is invoked by `Fiber.schedule`. Strictly speaking, you should use
294324
# it to create scheduled fibers, but it is not required in practice;
295325
# `Fiber.new` is usually sufficient.

0 commit comments

Comments
 (0)