Merge branch 'fix_forking'

mjsteger · mjsteger · commit 831766813fd9 · 2015-10-06T13:12:29.000-07:00
diff --git a/aws-flow/lib/aws/decider/task_poller.rb b/aws-flow/lib/aws/decider/task_poller.rb
@@ -52,7 +52,7 @@ def get_decision_task
         @domain.decision_tasks.poll_for_single_task(@task_list)
       end
 
-      def poll_and_process_single_task
+      def poll_and_process_single_task(opts={})
         # TODO waitIfSuspended
         begin
           @logger.debug "Polling for a new decision task of type #{@handler.workflow_definition_map.keys.map{ |x| "#{x.name} #{x.version}"} } on task_list: #{@task_list}"
@@ -96,6 +96,9 @@ def poll_and_process_single_task
           @logger.info Utilities.workflow_task_to_debug_string("Finished executing task", task, @task_list)
         rescue AWS::SimpleWorkflow::Errors::UnknownResourceFault => e
           @logger.error "Error in the poller, #{e.inspect}"
+        rescue Interrupt => e
+          @logger.error "Error in the poller, #{e.inspect}"
+          raise Interrupt
         rescue Exception => e
           @logger.error "Error in the poller, #{e.inspect}"
         end
@@ -356,7 +359,8 @@ def process_single_task(task)
       #   *Optional*. Whether to use forking to execute the task. On Windows,
       #   you should set this to `false`.
       #
-      def poll_and_process_single_task(use_forking = true)
+      def poll_and_process_single_task(opts = {})
+        use_forking = opts[:use_forking] || true
         @poll_semaphore ||= SuspendableSemaphore.new
         @poll_semaphore.acquire
         semaphore_needs_release = true
@@ -369,6 +373,10 @@ def poll_and_process_single_task(use_forking = true)
           if task
             @logger.info Utilities.activity_task_to_debug_string("Got activity task", task)
           end
+        rescue Interrupt => e
+          @poll_semaphore.release
+          @logger.error "Error in the poller, #{e.inspect}"
+          raise Interrupt
         rescue Exception => e
           @logger.error "Error in the poller, #{e.inspect}"
           @poll_semaphore.release
diff --git a/aws-flow/lib/aws/decider/worker.rb b/aws-flow/lib/aws/decider/worker.rb
@@ -43,17 +43,37 @@ def initialize(service, domain, task_list_to_poll, *args, &block)
         if args
           args.each { |klass_or_instance| add_implementation(klass_or_instance) }
         end
+        @aws_flow_signals = []
         @shutting_down = false
         %w{ TERM INT }.each do |signal|
           Signal.trap(signal) do
-            if @shutting_down
-              @executor.shutdown(0) if @executor
-              Kernel.exit! 1
-            else
-              @shutting_down = true
-              @shutdown_first_time_function.call if @shutdown_first_time_function
-            end
+            @aws_flow_signals << signal
+            raise Interrupt
+          end
+        end
+      end
+
+      def run_once(should_register = false, poller = nil)
+        register if should_register
+        poller ||= generate_poller
+        Kernel.exit if @shutting_down
+        poller.poll_and_process_single_task(@options)
+      end
+
+      def handle_signals
+        # This function itself needs to be able to handle interrupts, in case we get them in close succession
+        begin
+          return if @aws_flow_signals.empty?
+          if @shutting_down
+            @executor.shutdown(0) if @executor
+            Kernel.exit! 1
+          else
+            @shutting_down = true
+            @shutdown_first_time_function.call if @shutdown_first_time_function
           end
+        rescue Interrupt
+          @executor.shutdown(0) if @executor
+          Kernel.exit! 1
         end
       end
 
@@ -207,15 +227,15 @@ def register
       #   first. If {#register} was already called
       #   for this workflow worker, specify `false`.
       #
-      def start(should_register = true)
+      def start(should_register = true, poller = nil)
         # TODO check to make sure that the correct properties are set
         # TODO Register the domain if not already registered
         # TODO register types to poll
         # TODO Set up throttler
         # TODO Set up a timeout on the throttler correctly,
         # TODO Make this a generic poller, go to the right kind correctly
 
-        poller = WorkflowTaskPoller.new(
+        poller ||= WorkflowTaskPoller.new(
           @service,
           @domain,
           DecisionTaskHandler.new(@workflow_definition_map, @options),
@@ -226,10 +246,24 @@ def start(should_register = true)
         register if should_register
         @logger.debug "Starting an infinite loop to poll and process workflow tasks."
         loop do
-          run_once(false, poller)
+          begin
+            run_once(false, poller)
+          rescue Interrupt
+            handle_signals
+          end
         end
       end
 
+
+      def generate_poller
+        WorkflowTaskPoller.new(
+          @service,
+          @domain,
+          DecisionTaskHandler.new(@workflow_definition_map, @options),
+          @task_list,
+          @options
+        )
+      end
       # Starts the workflow and runs it once, with an optional
       # {WorkflowTaskPoller}.
       #
@@ -239,18 +273,7 @@ def start(should_register = true)
       #   An optional {WorkflowTaskPoller} to use.
       #
       def run_once(should_register = false, poller = nil)
-        register if should_register
-
-        poller = WorkflowTaskPoller.new(
-          @service,
-          @domain,
-          DecisionTaskHandler.new(@workflow_definition_map, @options),
-          @task_list,
-          @options
-        ) if poller.nil?
-
-        Kernel.exit if @shutting_down
-        poller.poll_and_process_single_task
+        super
       end
     end
 
@@ -391,7 +414,6 @@ def add_activities_implementation(class_or_instance)
         end
       end
 
-
       # Starts the activity that was added to the `ActivityWorker`.
       #
       # @param [true, false] should_register
@@ -412,10 +434,26 @@ def start(should_register = true)
 
         @logger.debug "Starting an infinite loop to poll and process activity tasks."
         loop do
-          run_once(false, poller)
+          begin
+            run_once(false, poller)
+          rescue Interrupt
+            handle_signals
+          end
+
         end
       end
 
+      def generate_poller
+        ActivityTaskPoller.new(
+          @service,
+          @domain,
+          @task_list,
+          @activity_definition_map,
+          @executor,
+          @options
+        )
+      end
+
       # Starts the activity that was added to the `ActivityWorker` and,
       # optionally, sets the {ActivityTaskPoller}.
       #
@@ -428,18 +466,7 @@ def start(should_register = true)
       #   {ActivityTaskPoller} will be created.
       #
       def run_once(should_register = true, poller = nil)
-        register if should_register
-        poller = ActivityTaskPoller.new(
-          @service,
-          @domain,
-          @task_list,
-          @activity_definition_map,
-          @executor,
-          @options
-        ) if poller.nil?
-
-        Kernel.exit if @shutting_down
-        poller.poll_and_process_single_task(@options.use_forking)
+        super
       end
     end
 
diff --git a/aws-flow/spec/aws/decider/unit/worker_spec.rb b/aws-flow/spec/aws/decider/unit/worker_spec.rb
@@ -39,6 +39,9 @@ def initialize(service, domain, task_list, forking_executor, *args, &block)
 end
 
 class FakeTaskPoller < WorkflowTaskPoller
+  def poll_and_process_single_task(opts={})
+    dumb_fib(5000)
+  end
   def get_decision_task
     nil
   end
@@ -134,48 +137,20 @@ class DefaultTasklistTestWorkflow
     workflow_worker = WorkflowWorker.new(service, domain, task_list)
     workflow_worker.add_workflow_implementation(TestWorkflow)
     pid = fork do
-      loop do
-        workflow_worker.run_once(true, FakeTaskPoller.new(service, domain, nil, task_list, nil))
-      end
+      workflow_worker.start(true, FakeTaskPoller.new(service, domain, nil, task_list, nil))
     end
+    sleep 1
     # Send an interrupt to the child process
     Process.kill("INT", pid)
     # Adding a sleep to let things get setup correctly (not ideal but going with
     # this for now)
-    sleep 5
+    sleep 2
     return_pid, status = Process.wait2(pid, Process::WNOHANG)
     Process.kill("KILL", pid) if return_pid.nil?
     return_pid.should_not be nil
     status.success?.should be_true
   end
 
-  it "will test whether WorkflowWorker dies cleanly when two interrupts are received" do
-    class FakeTaskPoller
-      def poll_and_process_single_task
-        dumb_fib(5000)
-      end
-    end
-    task_list = "TestWorkflow_tasklist"
-    service = FakeServiceClient.new
-    workflow_type_object = double("workflow_type", :name => "TestWorkflow.start", :start_execution => "" )
-    domain = FakeDomain.new(workflow_type_object)
-    workflow_worker = WorkflowWorker.new(service, domain, task_list)
-    workflow_worker.add_workflow_implementation(TestWorkflow)
-    pid = fork do
-      loop do
-        workflow_worker.run_once(true, FakeTaskPoller.new(service, domain, nil, task_list, nil))
-      end
-    end
-    # Send an interrupt to the child process
-    sleep 3
-    2.times { Process.kill("INT", pid); sleep 2 }
-    return_pid, status = Process.wait2(pid, Process::WNOHANG)
-
-    Process.kill("KILL", pid) if return_pid.nil?
-    return_pid.should_not be nil
-    status.success?.should be_false
-  end
-
 end
 
 describe ActivityWorker do
@@ -306,6 +281,7 @@ def activity_tasks
       sleep 30
     end
   end
+
   it "will test whether the ActivityWorker shuts down cleanly when an interrupt is received" do
 
     task_list = "TestWorkflow_tasklist"
@@ -319,20 +295,17 @@ def activity_tasks
     # handler to the process. When the process exits, the handler checks whether
     # the executor's internal is_shutdown variable is set correctly or not.
     pid = fork do
-      at_exit {
-        activity_worker.executor.is_shutdown.should == true
-      }
       activity_worker.start true
     end
     # Send an interrupt to the child process
+    sleep 1
     Process.kill("INT", pid)
     # Adding a sleep to let things get setup correctly (not ideal but going with
     # this for now)
-    sleep 5
+    sleep 2
     return_pid, status = Process.wait2(pid, Process::WNOHANG)
     Process.kill("KILL", pid) if return_pid.nil?
     return_pid.should_not be nil
-
     status.success?.should be_true
   end
 
@@ -354,15 +327,19 @@ def activity_tasks
     # create a child process to run that task. The task (dumb_fib) is
     # purposefully designed to be long running so that we can test our shutdown
     # scenario.
+
     pid = fork do
+      at_exit {
+        activity_worker.executor.is_shutdown.should == true
+      }
       activity_worker.executor.execute {
         dumb_fib(1000)
       }
       activity_worker.start true
     end
     # Adding a sleep to let things get setup correctly (not idea but going with
     # this for now)
-    sleep 3
+    sleep 2
     # Send 2 interrupts to the child process
     2.times { Process.kill("INT", pid); sleep 3 }
     status = Process.waitall
