From b65b27b02346be526bb4da2b70282c69050380aa Mon Sep 17 00:00:00 2001
From: Jeremy Mickelson <Jeremy.Mickelson@gmail.com>
Date: Mon, 22 Aug 2016 13:39:48 -0600
Subject: [PATCH] Make sure and let running polling tasks finish before
 shutting down

---
 aws-flow/lib/aws/decider/worker.rb | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/aws-flow/lib/aws/decider/worker.rb b/aws-flow/lib/aws/decider/worker.rb
index 5862d06..51e3bb1 100644
--- a/aws-flow/lib/aws/decider/worker.rb
+++ b/aws-flow/lib/aws/decider/worker.rb
@@ -51,7 +51,6 @@ def initialize(service, domain, task_list_to_poll, *args, &block)
               Kernel.exit! 1
             else
               @shutting_down = true
-              @shutdown_first_time_function.call if @shutdown_first_time_function
             end
           end
         end
@@ -304,10 +303,6 @@ def initialize(service, domain, task_list, *args, &block)
           :logger => @logger
         )
 
-        @shutdown_first_time_function = lambda do
-          @executor.shutdown Float::INFINITY
-          Kernel.exit
-        end
         super(service, domain, task_list, *args)
       end
 
@@ -438,7 +433,11 @@ def run_once(should_register = true, poller = nil)
           @options
         ) if poller.nil?
 
-        Kernel.exit if @shutting_down
+        if @shutting_down
+          @executor.shutdown Float::INFINITY
+          Kernel.exit
+        end
+
         poller.poll_and_process_single_task(@options.use_forking)
       end
     end