Preparing to fix #1812 - Moving locking code to a module

This moves the locking code out of Puppet::Util into a
separate module, to make the code cleaner.

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

lib/puppet/util.rb

@@ -10,6 +10,8 @@ module Puppet
 module Util
     require 'benchmark'
     
+    # These are all for backward compatibility -- these are methods that used
+    # to be in Puppet::Util but have been moved into external modules.
     require 'puppet/util/posix'
     extend Puppet::Util::POSIX
 
@@ -66,40 +68,6 @@ module Util
         end
     end
 
-    # Create a shared lock for reading
-    def self.readlock(file)
-        self.sync(file).synchronize(Sync::SH) do
-            File.open(file) { |f|
-                f.lock_shared { |lf| yield lf }
-            }
-        end
-    end
-
-    # Create an exclusive lock for writing, and do the writing in a
-    # tmp file.
-    def self.writelock(file, mode = 0600)
-        tmpfile = file + ".tmp"
-        unless FileTest.directory?(File.dirname(tmpfile))
-            raise Puppet::DevError, "Cannot create %s; directory %s does not exist" %
-                [file, File.dirname(file)]
-        end
-        self.sync(file).synchronize(Sync::EX) do
-            File.open(file, "w", mode) do |rf|
-                rf.lock_exclusive do |lrf|
-                    File.open(tmpfile, "w", mode) do |tf|
-                        yield tf
-                    end
-                    begin
-                        File.rename(tmpfile, file)
-                    rescue => detail
-                        Puppet.err "Could not rename %s to %s: %s" %
-                            [file, tmpfile, detail]
-                    end
-                end
-            end
-        end
-    end
-
     # Create instance methods for each of the log levels.  This allows
     # the messages to be a little richer.  Most classes will be calling this
     # method.

lib/puppet/util/file_locking.rb

@@ -0,0 +1,47 @@
+require 'puppet/util'
+
+module Puppet::Util::FileLocking
+    module_function
+
+    # Create a shared lock for reading
+    def readlock(file)
+        Puppet::Util.sync(file).synchronize(Sync::SH) do
+            File.open(file) { |f|
+                f.lock_shared { |lf| yield lf }
+            }
+        end
+    end
+
+    # Create an exclusive lock for writing, and do the writing in a
+    # tmp file.
+    def writelock(file, mode = nil)
+        unless FileTest.directory?(File.dirname(file))
+            raise Puppet::DevError, "Cannot create %s; directory %s does not exist" % [file, File.dirname(file)]
+        end
+        tmpfile = file + ".tmp"
+
+        unless mode
+            begin
+                mode = File.stat(file).mode
+            rescue
+                mode = 0600
+            end
+        end
+
+        Puppet::Util.sync(file).synchronize(Sync::EX) do
+            File.open(file, "w", mode) do |rf|
+                rf.lock_exclusive do |lrf|
+                    File.open(tmpfile, "w", mode) do |tf|
+                        yield tf
+                    end
+                    begin
+                        File.rename(tmpfile, file)
+                    rescue => detail
+                        File.unlink(tmpfile) if File.exist?(tmpfile)
+                        raise Puppet::Error, "Could not rename %s to %s: %s; file %s was unchanged" % [file, tmpfile, Thread.current.object_id, detail, file]
+                    end
+                end
+            end
+        end
+    end
+end

lib/puppet/util/storage.rb

@@ -1,6 +1,8 @@
 require 'yaml'
 require 'sync'
 
+require 'puppet/util/file_locking'
+
 # a class for storing state
 class Puppet::Util::Storage
     include Singleton
@@ -59,7 +61,7 @@ class Puppet::Util::Storage
             return
         end
         Puppet::Util.benchmark(:debug, "Loaded state") do
-            Puppet::Util.readlock(Puppet[:statefile]) do |file|
+            Puppet::Util::FileLocking.readlock(Puppet[:statefile]) do |file|
                 begin
                     @@state = YAML.load(file)
                 rescue => detail
@@ -97,7 +99,7 @@ class Puppet::Util::Storage
         end
 
         Puppet::Util.benchmark(:debug, "Stored state") do
-            Puppet::Util.writelock(Puppet[:statefile], 0660) do |file|
+            Puppet::Util::FileLocking.writelock(Puppet[:statefile], 0660) do |file|
                 file.print YAML.dump(@@state)
             end
         end

spec/integration/util/file_locking.rb

@@ -0,0 +1,36 @@
+#!/usr/bin/env ruby
+
+Dir.chdir(File.dirname(__FILE__)) { (s = lambda { |f| File.exist?(f) ? require(f) : Dir.chdir("..") { s.call(f) } }).call("spec/spec_helper.rb") }
+
+require 'puppet/util/file_locking'
+
+describe Puppet::Util::FileLocking do
+    it "should be able to keep file corruption from happening when there are multiple writers" do
+        file = Tempfile.new("puppetspec")
+        file.close!()
+        file = file.path
+        File.open(file, "w") { |f| f.puts "starting" }
+
+        value = {:a => :b}
+        threads = []
+        sync = Sync.new
+        9.times { |a|
+            threads << Thread.new {
+                9.times { |b|
+                    sync.synchronize(Sync::SH) {
+                        Puppet::Util::FileLocking.readlock(file) { |f|
+                            f.read
+                        }
+                    }
+                    sleep 0.01
+                    sync.synchronize(Sync::EX) {
+                        Puppet::Util::FileLocking.writelock(file) { |f|
+                            f.puts "%s %s" % [a, b]
+                        }
+                    }
+                }
+            }
+        }
+        threads.each { |th| th.join }
+    end
+end

spec/unit/util/file_locking.rb

@@ -0,0 +1,146 @@
+#!/usr/bin/env ruby
+
+Dir.chdir(File.dirname(__FILE__)) { (s = lambda { |f| File.exist?(f) ? require(f) : Dir.chdir("..") { s.call(f) } }).call("spec/spec_helper.rb") }
+
+require 'puppet/util/file_locking'
+
+class FileLocker
+    include Puppet::Util::FileLocking
+end
+
+describe Puppet::Util::FileLocking do
+    it "should have a module method for getting a read lock on files" do
+        Puppet::Util::FileLocking.should respond_to(:readlock)
+    end
+
+    it "should have a module method for getting a write lock on files" do
+        Puppet::Util::FileLocking.should respond_to(:writelock)
+    end
+
+    it "should have an instance method for getting a read lock on files" do
+        FileLocker.new.private_methods.should be_include("readlock")
+    end
+
+    it "should have an instance method for getting a write lock on files" do
+        FileLocker.new.private_methods.should be_include("writelock")
+    end
+
+    describe "when acquiring a read lock" do
+        it "should use a global shared mutex" do
+            sync = mock 'sync'
+            sync.expects(:synchronize).with(Sync::SH)
+            Puppet::Util.expects(:sync).with("/file").returns sync
+
+            Puppet::Util::FileLocking.readlock '/file'
+        end
+
+        it "should use a shared lock on the file" do
+            sync = mock 'sync'
+            sync.expects(:synchronize).yields
+            Puppet::Util.expects(:sync).with("/file").returns sync
+
+            fh = mock 'filehandle'
+            File.expects(:open).with("/file").yields fh
+            fh.expects(:lock_shared).yields "locked_fh"
+
+            result = nil
+            Puppet::Util::FileLocking.readlock('/file') { |l| result = l }
+            result.should == "locked_fh"
+        end
+    end
+
+    describe "when acquiring a write lock" do
+        before do
+            @sync = mock 'sync'
+            Puppet::Util.stubs(:sync).returns @sync
+            @sync.stubs(:synchronize).yields
+        end
+
+        it "should fail if the parent directory does not exist" do
+            FileTest.expects(:directory?).with("/my/dir").returns false
+
+            lambda { Puppet::Util::FileLocking.writelock('/my/dir/file') }.should raise_error(Puppet::DevError)
+        end
+
+        it "should use a global exclusive mutex" do
+            sync = mock 'sync'
+            sync.expects(:synchronize).with(Sync::EX)
+            Puppet::Util.expects(:sync).with("/file").returns sync
+
+            Puppet::Util::FileLocking.writelock '/file'
+        end
+
+        it "should use any specified mode when opening the file" do
+            File.expects(:open).with("/file", "w", :mymode)
+
+            Puppet::Util::FileLocking.writelock('/file', :mymode)
+        end
+
+        it "should use the mode of the existing file if no mode is specified" do
+            File.expects(:stat).with("/file").returns(mock("stat", :mode => 0755))
+            File.expects(:open).with("/file", "w", 0755)
+
+            Puppet::Util::FileLocking.writelock('/file')
+        end
+
+        it "should use 0600 as the mode if no mode is specified and the file does not exist" do
+            File.expects(:stat).raises(Errno::ENOENT)
+            File.expects(:open).with("/file", "w", 0600)
+
+            Puppet::Util::FileLocking.writelock('/file')
+        end
+
+        it "should create an exclusive file lock" do
+            fh = mock 'fh'
+            File.expects(:open).yields fh
+            fh.expects(:lock_exclusive)
+
+            Puppet::Util::FileLocking.writelock('/file')
+        end
+
+        it "should write to a temporary file" do
+            fh = mock 'fh'
+            File.expects(:open).yields fh
+
+            lfh = mock 'locked_filehandle'
+            fh.expects(:lock_exclusive).yields lfh
+
+            tf = mock 'tmp_filehandle'
+            File.expects(:open).with { |path, *args| path == "/file.tmp" }.yields tf
+
+            result = nil
+            File.stubs(:rename)
+            Puppet::Util::FileLocking.writelock('/file') { |f| result = f }
+            result.should equal(tf)
+        end
+
+        it "should rename the temporary file to the normal file" do
+            fh = stub 'fh'
+            fh.stubs(:lock_exclusive).yields fh
+            File.stubs(:open).yields fh
+
+            File.expects(:rename).with("/file.tmp", "/file")
+            Puppet::Util::FileLocking.writelock('/file') { |f| }
+        end
+
+        it "should fail if it cannot rename the file" do
+            fh = stub 'fh'
+            fh.stubs(:lock_exclusive).yields fh
+            File.stubs(:open).yields fh
+
+            File.expects(:rename).with("/file.tmp", "/file").raises(RuntimeError)
+            lambda { Puppet::Util::FileLocking.writelock('/file') { |f| } }.should raise_error(Puppet::Error)
+        end
+
+        it "should remove the temporary file if the rename fails" do
+            fh = stub 'fh'
+            fh.stubs(:lock_exclusive).yields fh
+            File.stubs(:open).yields fh
+
+            File.expects(:rename).with("/file.tmp", "/file").raises(RuntimeError)
+            File.expects(:exist?).with("/file.tmp").returns true
+            File.expects(:unlink).with("/file.tmp")
+            lambda { Puppet::Util::FileLocking.writelock('/file') { |f| } }.should raise_error(Puppet::Error)
+        end
+    end
+end

spec/unit/util/storage.rb

@@ -181,7 +181,7 @@ describe Puppet::Util::Storage do
             end
 
             it "should fail gracefully on load() if it cannot get a read lock on the state file" do
-                Puppet::Util.expects(:readlock).yields(false)
+                Puppet::Util::FileLocking.expects(:readlock).yields(false)
                 test_yaml = {'File["/yayness"]'=>{"name"=>{:a=>:b,:c=>:d}}}
                 YAML.expects(:load).returns(test_yaml)
 
@@ -223,7 +223,7 @@ describe Puppet::Util::Storage do
         end
 
         it "should raise an exception if it cannot get a write lock on the state file" do
-            Puppet::Util.expects(:writelock).yields(false)
+            Puppet::Util::FileLocking.expects(:writelock).yields(false)
             Puppet::Util::Storage.cache(:yayness)
 
             proc { Puppet::Util::Storage.store() }.should raise_error()

test/util/utiltest.rb

@@ -8,39 +8,6 @@ require 'mocha'
 class TestPuppetUtil < Test::Unit::TestCase
     include PuppetTest
 
-    # we're getting corrupt files, probably because multiple processes
-    # are reading or writing the file at once
-    # so we need to test that
-    def test_multiwrite
-        file = tempfile()
-        File.open(file, "w") { |f| f.puts "starting" }
-
-        value = {:a => :b}
-        threads = []
-        sync = Sync.new
-        9.times { |a|
-            threads << Thread.new {
-                9.times { |b|
-                    assert_nothing_raised {
-                        sync.synchronize(Sync::SH) {
-                            Puppet::Util.readlock(file) { |f|
-                                f.read
-                            }
-                        }
-                        sleep 0.01
-                        sync.synchronize(Sync::EX) {
-                            Puppet::Util.writelock(file) { |f|
-                                f.puts "%s %s" % [a, b]
-                            }
-                        }
-                    }
-                }
-            }
-        }
-        threads.each { |th| th.join }
-    end
-
-
     def test_withumask
         oldmask = File.umask