Fix 2239 (step two): introduce Puppet::Application.controlled_run method to provide simple status-restricted execution of a passed in block; this can replace the process status checks and properly handle delayed restart behavior for Puppet::Agent.

2 files changed, 43 insertions, 3 deletions

diff --git a/lib/puppet/application.rb b/lib/puppet/application.rb
index 83f90a18d..a80f42289 100644
--- a/lib/puppet/application.rb
+++ b/lib/puppet/application.rb
@@ -160,6 +160,18 @@ class Puppet::Application
         def clear?
             run_status.nil?
         end
+
+        # Only executes the given block if the run status of Puppet::Application is clear (no restarts, stops,
+        # etc. requested).
+        # Upon block execution, checks the run status again; if a restart has been requested during the block's
+        # execution, then controlled_run will send a new HUP signal to the current process.
+        # Thus, long-running background processes can potentially finish their work before a restart.
+        def controlled_run(&block)
+            return unless clear?
+            result = block.call
+            Process.kill(:HUP, $$) if restart_requested?
+            result
+        end
     end
 
     attr_reader :options, :opt_parser
diff --git a/spec/unit/application.rb b/spec/unit/application.rb
index cc8f79174..87a90099a 100755
--- a/spec/unit/application.rb
+++ b/spec/unit/application.rb
@@ -125,14 +125,42 @@ describe Puppet::Application do
         end
     end
 
-    describe 'when working with class-level run status properties' do
-        it 'should set run status and predicate appropriately on stop!' do
+    describe 'when performing a controlled_run' do
+        it 'should not execute block if not :clear?' do
+            Puppet::Application.run_status = :stop_requested
+            target = mock 'target'
+            target.expects(:some_method).never
+            Puppet::Application.controlled_run do
+                target.some_method
+            end
         end
 
-        it 'should set run status and predicate appropriately on restart!' do
+        it 'should execute block if :clear?' do
+            Puppet::Application.run_status = nil
+            target = mock 'target'
+            target.expects(:some_method).once
+            Puppet::Application.controlled_run do
+                target.some_method
+            end
         end
 
+        it 'should signal process with HUP after block if restart requested during block execution' do
+            Puppet::Application.run_status = nil
+            target = mock 'target'
+            target.expects(:some_method).once
+            old_handler = trap('HUP') { target.some_method }
+            begin
+                Puppet::Application.controlled_run do
+                    Puppet::Application.run_status = :restart_requested
+                end
+            ensure
+                trap('HUP', old_handler)
+            end
+        end
 
+        after :each do
+            Puppet::Application.run_status = nil
+        end
     end
 
     describe "when parsing command-line options" do