Preserve worker pool control flow during cancellation

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, &current_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) {

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