From 8ff44ebba02a7228b9c46f2ea80381eae4745805 Mon Sep 17 00:00:00 2001 From: Samuel Williams Date: Thu, 8 Aug 2024 13:50:50 +1200 Subject: [PATCH] Ensure that defer_stop resets state. (#339) --- lib/async/task.rb | 14 +++++++++----- test/async/task.rb | 12 ++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/lib/async/task.rb b/lib/async/task.rb index 7b28c028..9bb6c60a 100644 --- a/lib/async/task.rb +++ b/lib/async/task.rb @@ -304,19 +304,23 @@ def defer_stop # - false: defer_stop has been called and we are not stopping. # - true: defer_stop has been called and we will stop when exiting the block. if @defer_stop.nil? - # If we are not deferring stop already, we can defer it now: - @defer_stop = false - begin + # If we are not deferring stop already, we can defer it now: + @defer_stop = false + yield rescue Stop # If we are exiting due to a stop, we shouldn't try to invoke stop again: @defer_stop = nil raise ensure + defer_stop = @defer_stop + + # We need to ensure the state is reset before we exit the block: + @defer_stop = nil + # If we were asked to stop, we should do so now: - if @defer_stop - @defer_stop = nil + if defer_stop raise Stop, "Stopping current task (was deferred)!" end end diff --git a/test/async/task.rb b/test/async/task.rb index e2e7f3eb..87dac949 100644 --- a/test/async/task.rb +++ b/test/async/task.rb @@ -909,6 +909,18 @@ def sleep_forever child_task.stop expect(child_task).to be(:stopped?) end + + it "can defer stop and exit normally" do + child_task = reactor.async do |task| + task.defer_stop do + # Nothing. + end + end + + reactor.run_once(0) + + expect(child_task.stop_deferred?).to be == nil + end end with "failing task" do