Moving the Agent locking code to a module.

Also cleaning up the lock usage by yielding to a block
when a lock is attained.

Signed-off-by: Luke Kanies <luke@madstop.com>

4 files changed, 150 insertions, 66 deletions

diff --git a/lib/puppet/agent.rb b/lib/puppet/agent.rb
index 5cb6f0019..0046ed317 100644
--- a/lib/puppet/agent.rb
+++ b/lib/puppet/agent.rb
@@ -7,9 +7,11 @@ require 'puppet/util'
 class Puppet::Agent
     require 'puppet/agent/fact_handler'
     require 'puppet/agent/plugin_handler'
+    require 'puppet/agent/locker'
 
     include Puppet::Agent::FactHandler
     include Puppet::Agent::PluginHandler
+    include Puppet::Agent::Locker
 
     # For benchmarking
     include Puppet::Util
@@ -52,17 +54,6 @@ class Puppet::Agent
             end
         end
     end
-
-    # Let the daemon run again, freely in the filesystem.  Frolick, little
-    # daemon!
-    def enable
-        lockfile.unlock(:anonymous => true)
-    end
-
-    # Stop the daemon from making any catalog runs.
-    def disable
-        lockfile.lock(:anonymous => true)
-    end
     
     # Just so we can specify that we are "the" instance.
     def initialize
@@ -150,24 +141,25 @@ class Puppet::Agent
         got_lock = false
         splay
         Puppet::Util.sync(:puppetrun).synchronize(Sync::EX) do
-            unless lockfile.lock
-                Puppet.notice "Lock file %s exists; skipping catalog run" % lockfile.lockfile
-                return
-            end
+            got_lock = lock do
+                unless catalog = retrieve_catalog
+                    Puppet.err "Could not retrieve catalog; skipping run"
+                    return
+                end
 
-            got_lock = true
-            unless catalog = retrieve_catalog
-                Puppet.err "Could not retrieve catalog; skipping run"
-                return
+                begin
+                    benchmark(:notice, "Finished catalog run") do
+                        catalog.apply(options)
+                    end
+                rescue => detail
+                    puts detail.backtrace if Puppet[:trace]
+                    Puppet.err "Failed to apply catalog: %s" % detail
+                end
             end
 
-            begin
-                benchmark(:notice, "Finished catalog run") do
-                    catalog.apply(options)
-                end
-            rescue => detail
-                puts detail.backtrace if Puppet[:trace]
-                Puppet.err "Failed to apply catalog: %s" % detail
+            unless got_lock
+                Puppet.notice "Lock file %s exists; skipping catalog run" % lockfile.lockfile
+                return
             end
 
             # Now close all of our existing http connections, since there's no
@@ -180,10 +172,6 @@ class Puppet::Agent
             # done with the run.
             Process.kill(:HUP, $$) if self.restart?
         end
-    ensure
-        # Just make sure we remove the lock file if we set it.
-        lockfile.unlock if got_lock and lockfile.locked?
-        clear()
     end
 
     def running?
@@ -253,14 +241,6 @@ class Puppet::Agent
         return textobjects
     end
 
-    def lockfile
-        unless defined?(@lockfile)
-            @lockfile = Puppet::Util::Pidlock.new(Puppet[:puppetdlockfile])
-        end
-
-        @lockfile
-    end
-
     def splayed?
         @splayed
     end
diff --git a/lib/puppet/agent/locker.rb b/lib/puppet/agent/locker.rb
new file mode 100644
index 000000000..03736b278
--- /dev/null
+++ b/lib/puppet/agent/locker.rb
@@ -0,0 +1,38 @@
+require 'puppet/util/pidlock'
+
+# Break out the code related to locking the agent.  This module is just
+# included into the agent, but having it here makes it easier to test.
+module Puppet::Agent::Locker
+    # Let the daemon run again, freely in the filesystem.
+    def enable
+        lockfile.unlock(:anonymous => true)
+    end
+
+    # Stop the daemon from making any catalog runs.
+    def disable
+        lockfile.lock(:anonymous => true)
+    end
+
+    # Yield if we get a lock, else do nothing.  Return
+    # true/false depending on whether we get the lock.
+    def lock
+        if lockfile.lock
+            begin
+                yield
+            ensure
+                lockfile.unlock
+            end
+            return true
+        else
+            return false
+        end
+    end
+
+    def lockfile
+        unless defined?(@lockfile)
+            @lockfile = Puppet::Util::Pidlock.new(Puppet[:puppetdlockfile])
+        end
+
+        @lockfile
+    end
+end
diff --git a/spec/unit/agent.rb b/spec/unit/agent.rb
index 922577cb6..e0ae3689e 100755
--- a/spec/unit/agent.rb
+++ b/spec/unit/agent.rb
@@ -14,6 +14,10 @@ describe Puppet::Agent do
     it "should include the Fact Handler module" do
         Puppet::Agent.ancestors.should be_include(Puppet::Agent::FactHandler)
     end
+
+    it "should include the Locker module" do
+        Puppet::Agent.ancestors.should be_include(Puppet::Agent::Locker)
+    end
 end
 
 describe Puppet::Agent, "when executing a catalog run" do
@@ -21,10 +25,7 @@ describe Puppet::Agent, "when executing a catalog run" do
         Puppet.settings.stubs(:use).returns(true)
         @agent = Puppet::Agent.new
         @agent.stubs(:splay)
-
-        @lockfile = stub 'lockfile', :lock => true, :locked? => false, :lockfile => "/my/lock/file", :unlock => true
-
-        @agent.stubs(:lockfile).returns @lockfile
+        @agent.stubs(:lock).yields.then.returns true
     end
 
     it "should splay" do
@@ -43,8 +44,8 @@ describe Puppet::Agent, "when executing a catalog run" do
         @agent.run
     end
 
-    it "should use a lockfile to make sure no other process is executing the catalog" do
-        @lockfile.expects(:lock).returns true
+    it "should retrieve the catalog if a lock is attained" do
+        @agent.expects(:lock).yields.then.returns true
 
         @agent.expects(:retrieve_catalog)
 
@@ -52,7 +53,7 @@ describe Puppet::Agent, "when executing a catalog run" do
     end
 
     it "should log and do nothing if the lock cannot be acquired" do
-        @lockfile.expects(:lock).returns false
+        @agent.expects(:lock).returns false
 
         @agent.expects(:retrieve_catalog).never
 
@@ -93,27 +94,6 @@ describe Puppet::Agent, "when executing a catalog run" do
         @agent.run
     end
 
-    it "should remove the lock file when done applying the catalog" do
-        catalog = stub 'catalog', :retrieval_duration= => nil, :apply => nil
-        @agent.expects(:retrieve_catalog).returns catalog
-
-        @lockfile.expects(:lock).returns true
-
-        @lockfile.expects(:unlock)
-
-        @agent.run
-    end
-
-    it "should remove the lock file even if there was an exception during the run" do
-        catalog = stub 'catalog', :retrieval_duration= => nil
-        @agent.expects(:retrieve_catalog).returns catalog
-
-        catalog.expects(:apply).raises "eh"
-
-        @lockfile.expects(:unlock)
-        @agent.run
-    end
-
     it "should HUP itself if it should be restarted" do
         catalog = stub 'catalog', :retrieval_duration= => nil, :apply => nil
         @agent.expects(:retrieve_catalog).returns catalog
diff --git a/spec/unit/agent/locker.rb b/spec/unit/agent/locker.rb
new file mode 100755
index 000000000..aae7c0c7b
--- /dev/null
+++ b/spec/unit/agent/locker.rb
@@ -0,0 +1,86 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../../spec_helper'
+require 'puppet/agent'
+require 'puppet/agent/locker'
+
+class LockerTester
+    include Puppet::Agent::Locker
+end
+
+describe Puppet::Agent::Locker do
+    before do
+        @locker = LockerTester.new
+    end
+
+    it "should use a Pidlock instance as its lockfile" do
+        @locker.lockfile.should be_instance_of(Puppet::Util::Pidlock)
+    end
+
+    it "should reuse the same lock file each time" do
+        @locker.lockfile.should equal(@locker.lockfile)
+    end
+
+    it "should use the lock file to anonymously lock the process when disabled" do
+        @locker.lockfile.expects(:lock).with(:anonymous => true)
+
+        @locker.disable
+    end
+
+    it "should use the lock file to anonymously unlock the process when enabled" do
+        @locker.lockfile.expects(:unlock).with(:anonymous => true)
+
+        @locker.enable
+    end
+
+    it "should have a method that yields when a lock is attained" do
+        @locker.lockfile.expects(:lock).returns true
+
+        yielded = false
+        @locker.lock do
+            yielded = true
+        end
+        yielded.should be_true
+    end
+
+    it "should return true when the lock method successfully locked" do
+        @locker.lockfile.expects(:lock).returns true
+
+        @locker.lock {}.should be_true
+    end
+
+    it "should return true when the lock method does not receive the lock" do
+        @locker.lockfile.expects(:lock).returns false
+
+        @locker.lock {}.should be_false
+    end
+
+    it "should not yield when the lock method does not receive the lock" do
+        @locker.lockfile.expects(:lock).returns false
+
+        yielded = false
+        @locker.lock { yielded = true }
+        yielded.should be_false
+    end
+
+    it "should not unlock when a lock was not received" do
+        @locker.lockfile.expects(:lock).returns false
+        @locker.lockfile.expects(:unlock).never
+
+        @locker.lock {}
+    end
+
+    it "should unlock after yielding upon obtaining a lock" do
+        @locker.lockfile.stubs(:lock).returns true
+        @locker.lockfile.expects(:unlock)
+
+        @locker.lock {}
+    end
+
+    it "should unlock after yielding upon obtaining a lock, even if the block throws an exception" do
+        @locker.lockfile.stubs(:lock).returns true
+        @locker.lockfile.expects(:unlock)
+
+        lambda { @locker.lock { raise "foo" } }.should raise_error(RuntimeError)
+    end
+end