lib/jflow/activity/worker.rb in jflow-0.2.5 vs lib/jflow/activity/worker.rb in jflow-0.2.6

@@ -8,37 +8,47 @@
         @domain = domain
         @tasklist = tasklist
       end
 
       def start!
-        loop do
+        while should_be_working?
           log "Polling for task on #{domain} - #{tasklist}"
           poll
         end
+        log "Thread is marked as exiting, stopping the poll"
       end
 
-      private
 
       def poll
+        Thread.current["state"] = :polling
         response = JFlow.configuration.swf_client.poll_for_activity_task(poll_params)
         if response.task_token
-          process_task(response)
+          log "Got task #{response.task_token}"
+          task = JFlow::Activity::Task.new(response)
+          if should_be_working?
+            process(task)
+          else
+            #The worker is shuting down, we don't want to start working on anything
+            #so we fail the task and let the decider queue it up for retry later
+            task.failed!(Exception.new("Worker is going down!"))
+          end
         else
           log "Got no task"
         end
       end
 
-      def process_task(response)
-        log "Got task #{response.task_token}"
-        task = JFlow::Activity::Task.new(response)
+      def process(task)
         begin
+          Thread.current["state"] = :working
           task.run!
         rescue => exception
           task.failed!(exception)
         end
       end
 
+      private
+
       def identity
         @identity ||= "#{`hostname`.chomp}-#{Thread.current.object_id}"
       end
 
       def log(str)
@@ -51,9 +61,13 @@
           task_list: {
             name: tasklist,
           },
           identity: identity,
         }
+      end
+
+      def should_be_working?
+        Thread.current["do_exit"] != true
       end
 
     end
   end
 end
\ No newline at end of file