diff --git a/ext/io/event/worker_pool.c b/ext/io/event/worker_pool.c index 71233c4e..c58ec825 100644 --- a/ext/io/event/worker_pool.c +++ b/ext/io/event/worker_pool.c @@ -306,15 +306,54 @@ static VALUE worker_pool_initialize(int argc, VALUE *argv, VALUE self) { return self; } -static VALUE worker_pool_work_begin(VALUE _work) { +static VALUE worker_pool_work_block(VALUE _work) { struct IO_Event_WorkerPool_Work *work = (void*)_work; - if (DEBUG) fprintf(stderr, "worker_pool_work_begin:rb_fiber_scheduler_block work=%p\n", work); + if (DEBUG) fprintf(stderr, "worker_pool_work_block:rb_fiber_scheduler_block work=%p\n", work); rb_fiber_scheduler_block(work->scheduler, work->blocker, Qnil); return Qnil; } +static VALUE worker_pool_work_wait(VALUE _work) { + struct IO_Event_WorkerPool_Work *work = (void*)_work; + + while (true) { + worker_pool_work_block(_work); + if (DEBUG) fprintf(stderr, "-- worker_pool_work_wait:work completed=%d\n", work->completed); + + if (work->completed) { + break; + } + + if (DEBUG) fprintf(stderr, "worker_pool_work_wait:rb_fiber_scheduler_blocking_operation_cancel\n"); + rb_fiber_scheduler_blocking_operation_cancel(work->blocking_operation); + } + + return Qtrue; +} + +static VALUE worker_pool_work_ensure(VALUE _work) { + struct IO_Event_WorkerPool_Work *work = (void*)_work; + + while (!work->completed) { + if (DEBUG) fprintf(stderr, "worker_pool_work_ensure:rb_fiber_scheduler_blocking_operation_cancel\n"); + rb_fiber_scheduler_blocking_operation_cancel(work->blocking_operation); + + int state = 0; + rb_protect(worker_pool_work_block, _work, &state); + if (DEBUG) fprintf(stderr, "-- worker_pool_work_ensure:work completed=%d, state=%d\n", work->completed, state); + + if (state) { + // Ignore errors raised while waiting for cancellation to complete. rb_ensure + // will restore and rethrow the original control-flow state after this returns. + rb_set_errinfo(Qnil); + } + } + + return Qnil; +} + // Ruby method to submit work and wait for completion static VALUE worker_pool_call(VALUE self, VALUE _blocking_operation) { struct IO_Event_WorkerPool *pool; @@ -357,38 +396,10 @@ static VALUE worker_pool_call(VALUE self, VALUE _blocking_operation) { pthread_cond_signal(&pool->work_available); pthread_mutex_unlock(&pool->mutex); - // Block the current fiber until work is completed: - int state = 0; - while (true) { - int current_state = 0; - rb_protect(worker_pool_work_begin, (VALUE)&work, ¤t_state); - if (DEBUG) fprintf(stderr, "-- worker_pool_call:work completed=%d, current_state=%d, state=%d\n", work.completed, current_state, state); - - // Store the first exception state: - if (!state) { - state = current_state; - } - - // If the work is still in the queue, we must wait for a worker to complete it (even if cancelled): - if (work.completed) { - // The work was completed, we can exit the loop: - break; - } else { - if (DEBUG) fprintf(stderr, "worker_pool_call:rb_fiber_scheduler_blocking_operation_cancel\n"); - // Ensure the blocking operation is cancelled: - rb_fiber_scheduler_blocking_operation_cancel(blocking_operation); - - // The work was not completed, we need to wait for it to be completed, so we go around the loop again. - } - } - - if (DEBUG) fprintf(stderr, "<- worker_pool_call:work completed=%d, state=%d\n", work.completed, state); - - if (state) { - rb_jump_tag(state); - } else { - return Qtrue; - } + // Block the current fiber until work is completed. If blocking exits via an + // exception or another non-local jump, ensure the work is cancelled and fully + // drained before Ruby restores and rethrows the original control-flow state. + return rb_ensure(worker_pool_work_wait, (VALUE)&work, worker_pool_work_ensure, (VALUE)&work); } static VALUE worker_pool_allocate(VALUE klass) { diff --git a/test/io/event/worker_pool.rb b/test/io/event/worker_pool.rb index 24cbcd44..8ac3543d 100644 --- a/test/io/event/worker_pool.rb +++ b/test/io/event/worker_pool.rb @@ -225,4 +225,103 @@ ) end end + + with "scheduler that rescues internally in block" do + # A scheduler whose block method raises on the first call + # and rescues on a subsequent call clearing ec->errinfo. + class RescuingBlockScheduler < IO::Event::TestScheduler + attr_reader :rescued + + def initialize(...) + super + @block_call_count = 0 + end + + def block(blocker, timeout = nil) + @block_call_count += 1 + + raise StandardError, "first block failure" if @block_call_count == 1 + + begin + raise StandardError, "internal block error" + rescue + @rescued = true + end + super + end + end + + it "preserves the original exception through block rescue" do + result = nil + rescued = false + + Thread.new do + scheduler = RescuingBlockScheduler.new + Fiber.set_scheduler(scheduler) + + Fiber.schedule do + result = IO::Event::WorkerPool.busy(duration: 2.0) + end + rescued = scheduler.rescued + ensure + Fiber.set_scheduler(nil) + scheduler&.close + end.value + + expect(result).to be_a(Hash) + expect(result[:result]).to be == :exception + expect(result[:exception]).to be_a(StandardError) + expect(result[:exception].message).to be == "first block failure" + expect(rescued).to be == true + end + end + + with "scheduler that throws from block" do + class ThrowingBlockScheduler < IO::Event::TestScheduler + attr_reader :rescued + + def initialize(...) + super + @block_call_count = 0 + end + + def block(...) + @block_call_count += 1 + + throw :worker_pool_escape, :expected if @block_call_count == 1 + + begin + raise StandardError, "internal block error" + rescue + @rescued = true + end + + super + end + end + + it "preserves non-exception control flow" do + result = nil + rescued = false + + Thread.new do + scheduler = ThrowingBlockScheduler.new + Fiber.set_scheduler(scheduler) + + Fiber.schedule do + result = catch(:worker_pool_escape) do + IO::Event::WorkerPool.busy(duration: 2.0) + :missing + end + end + rescued = scheduler.rescued + ensure + Fiber.set_scheduler(nil) + scheduler&.close + end.value + + expect(result).to be == :expected + expect(rescued).to be == true + end + end end