[Feature] Implement Worker cancellation on signals and graceful shutdown timeout (#103)

* Add Signal handling to Temporalio::Worker

* Implement activity cancellation on a shutdown after a graceful timeout

* Switch to using ActivityCancelled error for Cancellations

* Rename stop_on_signals to shutdown_signals

Author: antstorm

.rubocop.yml

@@ -76,5 +76,5 @@ Metrics/PerceivedComplexity:
 Naming/MemoizedInstanceVariableName:
   Enabled: false
 
-Name/VariableNumber:
+Naming/VariableNumber:
   EnforcedStyle: snake_case

README.md

@@ -154,12 +154,23 @@ Temporal::Worker.run(worker_1, worker_2, worker_3)
 Please note that similar to `Temporal::Worker#run`, `Temporal::Worker.run` accepts a block that
 behaves the same way — the workers will be shut down when the block finishes.
 
+You can also configure your worker to listen on process signals to initiate a shutdown:
+
+```ruby
+Temporal::Worker.run(worker_1, worker_2, worker_3, shutdown_signals: %w[INT TERM])
+```
+
 #### Worker Shutdown
 
 The `Temporal::Worker#run` (as well as `Temporal::Worker#shutdown`) invocation will wait on all
 activities to complete, so if a long-running activity does not at least respect cancellation, the
 shutdown may never complete.
 
+If a worker was initialized with a `graceful_shutdown_timeout` option then a cancellation will be
+issued for every running activity after the set timeout. The bahaviour is mostly identical to a
+server requested cancellation and should be handled accordingly. More on this in [Heartbeating and
+Cancellation](#heartbeating-and-cancellation).
+
 ### Activities
 
 #### Definition
@@ -206,7 +217,7 @@ persisted on the server for retrieval during activity retry. If an activity call
 return an array containing `123` and `456` on the next run.
 
 A cancellation is implemented using the `Thread#raise` method, which will raise a
-`Temporal::Error::CancelledError` during the execution of an activity. This means that your code
+`Temporal::Error::ActivityCancelled` during the execution of an activity. This means that your code
 might get interrupted at any point and never complete. In order to protect critical parts of your
 activities wrap them in `activity.shield`:
 
@@ -245,6 +256,9 @@ end
 For any long-running activity using this approach it is recommended to periodically check
 `activity.cancelled?` flag and respond accordingly.
 
+Please note that your activities can also get cancelled during a worker shutdown process ([if
+configured accordingly](#worker-shutdown)).
+
 
 ## Dev Setup
 

lib/temporalio/activity/context.rb

@@ -1,4 +1,5 @@
 require 'temporalio/error/failure'
+require 'temporalio/errors'
 
 module Temporalio
   class Activity
@@ -14,7 +15,7 @@ def initialize(info, heartbeat_proc, shielded: false)
         @thread = Thread.current
         @info = info
         @heartbeat_proc = heartbeat_proc
-        @cancelled = false
+        @pending_cancellation = nil
         @shielded = shielded
         @mutex = Mutex.new
       end
@@ -41,7 +42,7 @@ def heartbeat(*details)
       def shield(&block)
         # The whole activity is shielded, called from a nested shield
         #   or while handling a CancelledError (in a rescue or ensure blocks)
-        return block.call if @shielded || @cancelled
+        return block.call if @shielded || cancelled?
 
         if Thread.current != thread
           # TODO: Use logger instead when implemented
@@ -52,7 +53,8 @@ def shield(&block)
         mutex.synchronize do
           @shielded = true
           result = block.call
-          raise Temporalio::Error::CancelledError, 'Unhandled cancellation' if @cancelled
+          # RBS: StandardError fallback is only needed to satisfy steep - https://github.com/soutaro/steep/issues/477
+          raise @pending_cancellation || StandardError if cancelled?
 
           result
         ensure # runs while still holding the lock
@@ -64,19 +66,23 @@ def shield(&block)
       #
       # @return [Boolean] true if the activity has had a cancellation request, false otherwise.
       def cancelled?
-        @cancelled
+        !!@pending_cancellation
       end
 
-      # Cancel the running activity.
+      # Cancel the running activity by raising a provided error.
+      #
+      # @param reason [String] Reason for cancellation.
+      # @param by_request [Boolean] Cancellation requested by the server or not.
       #
       # @api private
-      def cancel
-        @cancelled = true
+      def cancel(reason, by_request: false)
+        error = Temporalio::Error::ActivityCancelled.new(reason, by_request)
+        @pending_cancellation = error
 
         locked = mutex.try_lock
         # @shielded inside the lock means the whole activity is shielded
         if locked && !@shielded
-          thread.raise(Temporalio::Error::CancelledError.new('Unhandled cancellation'))
+          thread.raise(error)
         end
       ensure
         # only release the lock if we locked it

lib/temporalio/errors.rb

@@ -23,5 +23,18 @@ def initialize(status)
     class Internal < Error; end
 
     class WorkerShutdown < Internal; end
+
+    # This error is used within the SDK to communicate Activity cancellations
+    # (and whether it was requested by server or not)
+    class ActivityCancelled < Internal
+      def initialize(reason, by_request)
+        super(reason)
+        @by_request = by_request
+      end
+
+      def by_request?
+        @by_request
+      end
+    end
   end
 end

lib/temporalio/worker.rb

@@ -17,11 +17,26 @@ class Worker
     # finished) and will raise if any of the workers raises a fatal error.
     #
     # @param workers [Array<Temporalio::Worker>] A list of the workers to be run.
+    # @param shutdown_signals [Array<String>] A list of process signals for the worker to stop on.
+    #   This argument can not be used with a custom block.
     #
     # @yield Optionally you can provide a block by the end of which all the workers will be shut
     #   down. Any errors raised from this block will be re-raised by this method.
-    def self.run(*workers, &block)
-      # TODO: Add signal handling
+    def self.run(*workers, shutdown_signals: [], &block)
+      unless shutdown_signals.empty?
+        if block
+          raise ArgumentError, 'Temporalio::Worker.run accepts :shutdown_signals or a block, but not both'
+        end
+
+        signal_queue = Queue.new
+
+        shutdown_signals.each do |signal|
+          Signal.trap(signal) { signal_queue.close }
+        end
+
+        block = -> { signal_queue.pop }
+      end
+
       Runner.new(*workers).run(&block)
     end
 
@@ -36,6 +51,9 @@ def self.run(*workers, &block)
     # @param activity_executor [ThreadPoolExecutor] Concurrent executor for all activities. Defaults
     #   to a {ThreadPoolExecutor} with `:max_concurrent_activities` available threads.
     # @param max_concurrent_activities [Integer] Number of concurrently running activities.
+    # @param graceful_shutdown_timeout [Integer] Amount of time (in seconds) activities are given
+    #   after a shutdown to complete before they are cancelled. A default value of `nil` means that
+    #   activities are never cancelled when handling a shutdown.
     #
     # @raise [ArgumentError] When no activities or workflows have been provided.
     def initialize(
@@ -45,7 +63,8 @@ def initialize(
       activities: [],
       data_converter: Temporalio::DataConverter.new,
       activity_executor: nil,
-      max_concurrent_activities: 100
+      max_concurrent_activities: 100,
+      graceful_shutdown_timeout: nil
     )
       # TODO: Add worker interceptors
       @started = false
@@ -59,13 +78,17 @@ def initialize(
         namespace,
         task_queue,
       )
-      @activity_worker = init_activity_worker(
-        task_queue,
-        @core_worker,
-        activities,
-        data_converter,
-        @activity_executor,
-      )
+      @activity_worker =
+        unless activities.empty?
+          Worker::ActivityWorker.new(
+            task_queue,
+            @core_worker,
+            activities,
+            data_converter,
+            @activity_executor,
+            graceful_shutdown_timeout,
+          )
+        end
       @workflow_worker = nil
 
       if !@activity_worker && !@workflow_worker
@@ -171,11 +194,5 @@ def running?
 
     attr_reader :mutex, :runtime, :activity_executor, :core_worker, :activity_worker,
                 :workflow_worker, :runner
-
-    def init_activity_worker(task_queue, core_worker, activities, data_converter, executor)
-      return if activities.empty?
-
-      Worker::ActivityWorker.new(task_queue, core_worker, activities, data_converter, executor)
-    end
   end
 end

lib/temporalio/worker/activity_runner.rb

@@ -1,6 +1,8 @@
 require 'google/protobuf/well_known_types'
 require 'temporalio/activity/context'
 require 'temporalio/activity/info'
+require 'temporalio/error/failure'
+require 'temporalio/errors'
 
 module Temporalio
   class Worker
@@ -27,11 +29,21 @@ def run
 
         converter.to_payload(result)
       rescue StandardError => e
+        # Temporal server ignores cancellation failures that were not requested by the server.
+        # However within the SDK cancellations are also used during the worker shutdown. In order
+        # to provide a seamless handling experience (same error raised within the Activity) we are
+        # using the ActivityCancelled error and then swapping it with a CancelledError here.
+        #
+        # In the future this will be handled by the SDK Core — https://github.com/temporalio/sdk-core/issues/461
+        if e.is_a?(Temporalio::Error::ActivityCancelled) && e.by_request?
+          e = Temporalio::Error::CancelledError.new(e.message)
+        end
+
         converter.to_failure(e)
       end
 
-      def cancel
-        context.cancel
+      def cancel(reason, by_request:)
+        context.cancel(reason, by_request: by_request)
       end
 
       private

lib/temporalio/worker/activity_worker.rb

@@ -7,12 +7,13 @@ module Temporalio
   class Worker
     # @api private
     class ActivityWorker
-      def initialize(task_queue, core_worker, activities, converter, executor)
+      def initialize(task_queue, core_worker, activities, converter, executor, graceful_timeout)
         @task_queue = task_queue
         @worker = SyncWorker.new(core_worker)
         @activities = prepare_activities(activities)
         @converter = converter
         @executor = executor
+        @graceful_timeout = graceful_timeout
         @running_activities = {}
         @cancellations = []
         @drain_queue = Queue.new
@@ -37,8 +38,19 @@ def run(reactor)
       rescue Temporalio::Bridge::Error::WorkerShutdown
         # No need to re-raise this error, it's a part of a normal shutdown
       ensure
-        reactor.async do
+        reactor.async do |async_task|
+          cancelation_task =
+            if graceful_timeout
+              async_task.async do
+                sleep graceful_timeout
+                running_activities.each_value do |activity_runner|
+                  activity_runner.cancel('Worker is shutting down', by_request: false)
+                end
+              end
+            end
+
           outstanding_tasks.each(&:wait)
+          cancelation_task&.stop # all tasks completed, stop cancellations
           drain_queue.close
         end
       end
@@ -49,8 +61,8 @@ def drain
 
       private
 
-      attr_reader :task_queue, :worker, :activities, :converter, :executor, :running_activities,
-                  :cancellations, :drain_queue
+      attr_reader :task_queue, :worker, :activities, :converter, :executor, :graceful_timeout,
+                  :running_activities, :cancellations, :drain_queue
 
       def prepare_activities(activities)
         activities.each_with_object({}) do |activity, result|
@@ -119,7 +131,7 @@ def handle_cancel_activity(task_token, _cancel)
         end
 
         cancellations << task_token
-        runner&.cancel
+        runner&.cancel('Activity cancellation requested', by_request: true)
       end
     end
   end

lib/temporalio/worker/runner.rb

@@ -1,3 +1,5 @@
+require 'temporalio/errors'
+
 module Temporalio
   class Worker
     # A class used to manage the lifecycle of running any number of workers.

sig/async.rbs

@@ -8,6 +8,7 @@ module Async
   class Task
     def async: { (Async::Task) -> void } -> Async::Task
     def wait: -> void
+    def stop: -> void
   end
 end
 

sig/temporalio/activity/context.rbs

@@ -7,12 +7,12 @@ module Temporalio
       def heartbeat: (*untyped details) -> void
       def shield: { -> untyped } -> untyped
       def cancelled?: -> bool
-      def cancel: -> void
+      def cancel: (String reason, ?by_request: bool) -> void
 
       private
 
       @shielded: bool
-      @cancelled: bool
+      @pending_cancellation: Exception?
 
       attr_reader thread: Thread
       attr_reader heartbeat_proc: ^(*untyped) -> void