diff --git a/lib/async/task.rb b/lib/async/task.rb index 2b02f689..bc7bd00c 100644 --- a/lib/async/task.rb +++ b/lib/async/task.rb @@ -67,6 +67,8 @@ def initialize(parent = Task.current?, finished: nil, **options, &block) @status = :initialized @result = nil @finished = finished + + @defer_stop = nil end def reactor @@ -212,6 +214,13 @@ def stop(later = false) return stopped! end + # If we are deferring stop... + if @defer_stop == false + # Don't stop now... but update the state so we know we need to stop later. + @defer_stop = true + return false + end + # If the fiber is alive, we need to stop it: if @fiber&.alive? if self.current? @@ -239,6 +248,41 @@ def stop(later = false) end end + # Defer the handling of stop. During the execution of the given block, if a stop is requested, it will be deferred until the block exits. This is useful for ensuring graceful shutdown of servers and other long-running tasks. You should wrap the response handling code in a defer_stop block to ensure that the task is stopped when the response is complete but not before. + # + # You can nest calls to defer_stop, but the stop will only be deferred until the outermost block exits. + # + # If stop is invoked a second time, it will be immediately executed. + # + # @yields {} The block of code to execute. + def defer_stop + # Tri-state variable for controlling stop: + # - nil: defer_stop has not been called. + # - 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 + yield + rescue Stop + # If we are exiting due to a stop, we shouldn't try to invoke stop again: + @defer_stop = nil + raise + ensure + # If we were asked to stop, we should do so now: + if @defer_stop + @defer_stop = nil + self.stop + end + end + else + # If we are deferring stop already, entering it again is a no-op. + yield + end + end + # Lookup the {Task} for the current fiber. Raise `RuntimeError` if none is available. # @returns [Task] # @raises[RuntimeError] If task was not {set!} for the current fiber. diff --git a/test/async/task.rb b/test/async/task.rb index eb39522d..82d5c7b8 100644 --- a/test/async/task.rb +++ b/test/async/task.rb @@ -841,7 +841,7 @@ def sleep_forever it "can gets in a task" do IO.pipe do |input, output| - Async do + Async do Async do expect(input.gets).to be == "hello\n" end @@ -849,5 +849,62 @@ def sleep_forever end end end - + + with '#defer_stop' do + it "can defer stopping a task" do + child_task = reactor.async do |task| + task.defer_stop do + sleep + end + end + + reactor.run_once(0) + + child_task.stop + expect(child_task).to be(:running?) + + child_task.stop + expect(child_task).to be(:stopped?) + end + + it "will stop the task if it was deferred" do + condition = Async::Notification.new + + child_task = reactor.async do |task| + task.defer_stop do + condition.wait + end + end + + reactor.run_once(0) + + child_task.stop(true) + expect(child_task).to be(:running?) + + reactor.async do + condition.signal + end + + reactor.run_once(0) + expect(child_task).to be(:stopped?) + end + + it "can defer stop in a deferred stop" do + child_task = reactor.async do |task| + task.defer_stop do + task.defer_stop do + sleep + end + end + end + + reactor.run_once(0) + + child_task.stop + expect(child_task).to be(:running?) + + child_task.stop + expect(child_task).to be(:stopped?) + end + end end