Adding HUP and USR1 hooks

git-svn-id: https://reductivelabs.com/svn/puppet/trunk@1308 980ebf18-57e1-0310-9a29-db15c13687c0

4 files changed, 98 insertions, 6 deletions

diff --git a/lib/puppet.rb b/lib/puppet.rb
index 5ad4ae659..ef08e33fc 100644
--- a/lib/puppet.rb
+++ b/lib/puppet.rb
@@ -20,9 +20,13 @@ module Puppet
         return PUPPETVERSION
     end
 
-    # So we can monitor signals and such.
     class << self
+        # So we can monitor signals and such.
         include SignalObserver
+
+        # To keep a copy of arguments.  Set within Config#addargs, because I'm
+        # lazy.
+        attr_accessor :args
     end
 
     class Error < RuntimeError
@@ -320,6 +324,14 @@ module Puppet
         timer
     end
 
+    # Relaunch the executable.
+    def self.restart
+        command = $0 + " " + self.args.join(" ")
+        Puppet.notice "Restarting with '%s'" % command
+        Puppet.shutdown(false)
+        exec(command)
+    end
+
     # Trap a couple of the main signals.  This should probably be handled
     # in a way that anyone else can register callbacks for traps, but, eh.
     def self.settraps
@@ -329,19 +341,58 @@ module Puppet
                 Puppet.shutdown
             end
         end
+
+        # Handle restarting.
+        trap(:HUP) do
+            if client = @services.find { |s| s.is_a? Puppet::Client::MasterClient } and client.running?
+                client.restart
+            else
+                Puppet.restart
+            end
+        end
+
+        # Provide a hook for running clients where appropriate
+        trap(:USR1) do
+            done = 0
+            Puppet.notice "Caught USR1; triggering client run"
+            @services.find_all { |s| s.is_a? Puppet::Client }.each do |client|
+                if client.respond_to? :running?
+                    if client.running?
+                        Puppet.info "Ignoring running %s" % client.class
+                    else
+                        done += 1
+                        begin
+                            client.runnow
+                        rescue => detail
+                            Puppet.err "Could not run client: %s" % detail
+                        end
+                    end
+                else
+                    Puppet.info "Ignoring %s; cannot test whether it is running" %
+                        client.class
+                end
+            end
+
+            unless done > 0
+                Puppet.notice "No clients were run"
+            end
+        end
     end
 
     # Shutdown our server process, meaning stop all services and all threads.
     # Optionally, exit.
     def self.shutdown(leave = true)
+        Puppet.notice "Shutting down"
         # Unmonitor our timers
         defined? @timers and @timers.each do |timer|
             EventLoop.current.ignore_timer timer
         end
 
-        if EventLoop.current.running?
-            EventLoop.current.quit
-        end
+        # This seems to exit the process, although I can't find where it does
+        # so.  Leaving it out doesn't seem to hurt anything.
+        #if EventLoop.current.running?
+        #    EventLoop.current.quit
+        #end
 
         # Stop our services
         defined? @services and @services.each do |svc|
@@ -387,11 +438,17 @@ module Puppet
                     if Puppet[:debug]
                         puts detail.backtrace
                     end
+                    @services.delete svc
                     Puppet.err "Could not start %s: %s" % [svc.class, detail]
                 end
             end
         end
 
+        unless @services.length > 0
+            Puppet.notice "No remaining services; exiting"
+            exit(1)
+        end
+
         if block
             EventLoop.current.run
         end
diff --git a/lib/puppet/client.rb b/lib/puppet/client.rb
index bb1572063..35501acc2 100644
--- a/lib/puppet/client.rb
+++ b/lib/puppet/client.rb
@@ -135,8 +135,9 @@ module Puppet
             else
                 self.stopping = true
                 rmpidfile()
-                Puppet::Storage.store
-                exit
+                if self.running?
+                    Puppet::Storage.store
+                end
             end
         end
 
diff --git a/lib/puppet/client/master.rb b/lib/puppet/client/master.rb
index 70037729f..cd9bb0802 100644
--- a/lib/puppet/client/master.rb
+++ b/lib/puppet/client/master.rb
@@ -309,6 +309,7 @@ class Puppet::Client::MasterClient < Puppet::Client
         super
 
         self.class.instance = self
+        @running = false
     end
 
     # Make sure only one client runs at a time, and make sure only one thread
@@ -342,6 +343,23 @@ class Puppet::Client::MasterClient < Puppet::Client
         end
     end
 
+    # Mark that we should restart.  The Puppet module checks whether we're running,
+    # so this only gets called if we're in the middle of a run.
+    def restart
+        # If we're currently running, then just mark for later
+        Puppet.notice "Received signal to restart; waiting until run is complete"
+        @restart = true
+    end
+
+    # Should we restart?
+    def restart?
+        if defined? @restart
+            @restart
+        else
+            false
+        end
+    end
+
     # Retrieve the cached config
     def retrievecache
         if FileTest.exists?(self.cachefile)
@@ -362,6 +380,7 @@ class Puppet::Client::MasterClient < Puppet::Client
                 Puppet[:puppetdlockfile]
         else
             lock do
+                @running = true
                 self.getconfig
 
                 if defined? @objects and @objects
@@ -372,10 +391,21 @@ class Puppet::Client::MasterClient < Puppet::Client
                         self.apply(tags, ignoreschedules)
                     end
                 end
+                @running = false
+            end
+
+            # Did we get HUPped during the run?  If so, then restart now that we're
+            # done with the run.
+            if self.restart?
+                Process.kill(:HUP, $$)
             end
         end
     end
 
+    def running?
+        @running
+    end
+
     # Store the classes in the classfile, but only if we're not local.
     def setclasses(ary)
         if @local
diff --git a/lib/puppet/config.rb b/lib/puppet/config.rb
index bcfdabd90..3c76a623f 100644
--- a/lib/puppet/config.rb
+++ b/lib/puppet/config.rb
@@ -50,6 +50,10 @@ class Config
     # understand, and add them to the passed option list.
     def addargs(options)
         require 'getoptlong'
+
+        # Hackish, but acceptable.  Copy the current ARGV for restarting.
+        Puppet.args = ARGV.dup
+
         # Add all of the config parameters as valid options.
         self.each { |param, obj|
             if self.boolean?(param)