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

@@ -17,15 +17,15 @@
         log "Thread is marked as exiting, stopping the poll"
       end
 
 
       def poll
-        Thread.current["state"] = :polling
+        Thread.current.set_state(:polling)
         response = JFlow.configuration.swf_client.poll_for_activity_task(poll_params)
         if response.task_token
-          log "Got task #{response.task_token}"
           task = JFlow::Activity::Task.new(response)
+          log "Got task #{task.workflow_id}-#{task.run_id}"
           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
@@ -36,11 +36,11 @@
         end
       end
 
       def process(task)
         begin
-          Thread.current["state"] = :working
+          Thread.current.set_state(:working)
           task.run!
         rescue => exception
           task.failed!(exception)
         end
       end
@@ -64,10 +64,10 @@
           identity: identity,
         }
       end
 
       def should_be_working?
-        Thread.current["do_exit"] != true
+        !Thread.current.marked_for_shutdown?
       end
 
     end
   end
 end
\ No newline at end of file