Fixing #447 - filebuckets now use deeply nested paths

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

3 files changed, 180 insertions, 68 deletions

diff --git a/CHANGELOG b/CHANGELOG
index fffb9697f..ff0156b6e 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,4 +1,8 @@
 0.22.2 (grover)
+    FileBuckets now use a deeply nested structure for storing files, so
+    you do not end up with hundreds or thousands of files in the same
+    directory. (#447)
+
     Facts are now cached in the state file, and when they change the configuration
     is always recompiled. (#519)
 
diff --git a/lib/puppet/network/handler/filebucket.rb b/lib/puppet/network/handler/filebucket.rb
index 653d566b4..96b9a1e9a 100755
--- a/lib/puppet/network/handler/filebucket.rb
+++ b/lib/puppet/network/handler/filebucket.rb
@@ -1,16 +1,12 @@
-#--------------------
-# accept and serve files
-
-
-require 'webrick'
-require 'xmlrpc/server'
-require 'xmlrpc/client'
-require 'facter'
+require 'fileutils'
 require 'digest/md5'
 require 'puppet/external/base64'
 
-class Puppet::Network::Handler
+class Puppet::Network::Handler # :nodoc:
     class BucketError < RuntimeError; end
+    # Accept files and store them by md5 sum, returning the md5 sum back
+    # to the client.  Alternatively, accept an md5 sum and return the
+    # associated content.
     class FileBucket < Handler
         @interface = XMLRPC::Service::Interface.new("puppetbucket") { |iface|
             iface.add_method("string addfile(string, string)")
@@ -21,7 +17,7 @@ class Puppet::Network::Handler
         attr_reader :name, :path
 
         # this doesn't work for relative paths
-        def self.paths(base,md5)
+        def self.oldpaths(base,md5)
             return [
                 File.join(base, md5),
                 File.join(base, md5, "contents"),
@@ -29,12 +25,29 @@ class Puppet::Network::Handler
             ]
         end
 
+        # this doesn't work for relative paths
+        def self.paths(base,md5)
+            dir = File.join(md5[0..7].split(""))
+            basedir = File.join(base, dir)
+            return [
+                basedir,
+                File.join(basedir, "contents"),
+                File.join(basedir, "paths")
+            ]
+        end
+
+        # Should we check each file as it comes in to make sure the md5
+        # sums match?  Defaults to false.
+        def conflict_check?
+            @confictchk
+        end
+
         def initialize(hash)
             if hash.include?(:ConflictCheck)
                 @conflictchk = hash[:ConflictCheck]
                 hash.delete(:ConflictCheck)
             else
-                @conflictchk = true
+                @conflictchk = false
             end
 
             if hash.include?(:Path)
@@ -53,7 +66,8 @@ class Puppet::Network::Handler
             @name = "Filebucket[#{@path}]"
         end
 
-        # accept a file from a client
+        # Accept a file from a client and store it by md5 sum, returning
+        # the sum.
         def addfile(contents, path, client = nil, clientip = nil)
             if client
                 contents = Base64.decode64(contents)
@@ -62,71 +76,48 @@ class Puppet::Network::Handler
 
             bpath, bfile, pathpath = FileBucket.paths(@path,md5)
 
-            # if it's a new directory...
-            if Puppet.recmkdir(bpath)
-                msg = "Adding %s(%s)" % [path, md5]
-                msg += " from #{client}" if client
-                self.info msg
-                # ...then just create the file
-                File.open(bfile, File::WRONLY|File::CREAT, 0440) { |of|
-                    of.print contents
-                }
-            else # if the dir already existed...
-                # ...we need to verify that the contents match the existing file
-                if @conflictchk
-                    unless FileTest.exists?(bfile)
-                        raise(BucketError,
-                            "No file at %s for sum %s" % [bfile,md5], caller)
-                    end
-
-                    curfile = File.read(bfile)
-
-                    # If the contents don't match, then we've found a conflict.
-                    # Unlikely, but quite bad.
-                    if curfile != contents
-                        raise(BucketError,
-                            "Got passed new contents for sum %s" % md5, caller)
-                    else
-                        msg = "Got duplicate %s(%s)" % [path, md5]
-                        msg += " from #{client}" if client
-                        self.info msg
-                    end
+            # If the file already exists, just return the md5 sum.
+            if FileTest.exists?(bfile)
+                # If verification is enabled, then make sure the text matches.
+                if conflict_check?
+                    verify(contents, md5, bfile)
                 end
+                return md5
             end
 
-            contents = ""
-
-            # in either case, add the passed path to the list of paths
-            paths = nil
-            addpath = false
-            if FileTest.exists?(pathpath)
-                File.open(pathpath) { |of|
-                    paths = of.readlines.collect { |l| l.chomp }
-                }
-
-                # unless our path is already there...
-                unless paths.include?(path)
-                    addpath = true
+            # Make the directories if necessary.
+            unless FileTest.directory?(bpath)
+                Puppet::Util.withumask(0007) do
+                    FileUtils.mkdir_p(bpath)
                 end
-            else
-                addpath = true
             end
 
-            # if it's a new file, or if our path isn't in the file yet, add it
-            if addpath
-                File.open(pathpath, File::WRONLY|File::CREAT|File::APPEND) { |of|
-                    of.puts path
-                }
-            end
+            # Write the file to disk.
+            msg = "Adding %s(%s)" % [path, md5]
+            msg += " from #{client}" if client
+            self.info msg
+
+            # ...then just create the file
+            File.open(bfile, File::WRONLY|File::CREAT, 0440) { |of|
+                of.print contents
+            }
+
+            # Write the path to the paths file.
+            add_path(path, pathpath)
 
             return md5
         end
 
+        # Return the contents associated with a given md5 sum.
         def getfile(md5, client = nil, clientip = nil)
             bpath, bfile, bpaths = FileBucket.paths(@path,md5)
 
             unless FileTest.exists?(bfile)
-                return false
+                # Try the old flat style.
+                bpath, bfile, bpaths = FileBucket.oldpaths(@path,md5)
+                unless FileTest.exists?(bfile)
+                    return false
+                end
             end
 
             contents = nil
@@ -144,6 +135,39 @@ class Puppet::Network::Handler
         def to_s
             self.name
         end
+
+        private
+
+        # Add our path to the paths file if necessary.
+        def add_path(path, file)
+            if FileTest.exists?(file)
+                File.open(file) { |of|
+                    return if of.readlines.collect { |l| l.chomp }.include?(path)
+                }
+            end
+
+            # if it's a new file, or if our path isn't in the file yet, add it
+            File.open(file, File::WRONLY|File::CREAT|File::APPEND) { |of|
+                of.puts path
+            }
+        end
+
+        # If conflict_check is enabled, verify that the passed text is
+        # the same as the text in our file.
+        def verify(content, md5, bfile)
+            curfile = File.read(bfile)
+
+            # If the contents don't match, then we've found a conflict.
+            # Unlikely, but quite bad.
+            if curfile != contents
+                raise(BucketError,
+                    "Got passed new contents for sum %s" % md5, caller)
+            else
+                msg = "Got duplicate %s(%s)" % [path, md5]
+                msg += " from #{client}" if client
+                self.info msg
+            end
+        end
     end
 end
 
diff --git a/test/network/handler/bucket.rb b/test/network/handler/bucket.rb
index e1ba41498..9c68afd29 100755
--- a/test/network/handler/bucket.rb
+++ b/test/network/handler/bucket.rb
@@ -148,7 +148,7 @@ class TestBucket < Test::Unit::TestCase
         server = nil
         assert_nothing_raised {
             server = Puppet::Network::Handler.filebucket.new(
-                :Bucket => @bucket
+                :Path => @bucket
             )
         }
 
@@ -187,7 +187,7 @@ class TestBucket < Test::Unit::TestCase
         threads = []
         assert_nothing_raised {
             bucket = Puppet::Network::Handler.filebucket.new(
-                :Bucket => @bucket
+                :Path => @bucket
             )
         }
 
@@ -214,7 +214,7 @@ class TestBucket < Test::Unit::TestCase
         client = nil
         port = Puppet[:masterport]
 
-        pid = mkserver(:CA => {}, :FileBucket => { :Bucket => @bucket})
+        pid = mkserver(:CA => {}, :FileBucket => { :Path => @bucket})
 
         assert_nothing_raised {
             client = Puppet::Network::Client.dipper.new(
@@ -235,7 +235,7 @@ class TestBucket < Test::Unit::TestCase
         bucket = nil
         assert_nothing_raised {
             bucket = Puppet::Network::Handler.filebucket.new(
-                :Bucket => @bucket
+                :Path => @bucket
             )
         }
 
@@ -247,12 +247,96 @@ class TestBucket < Test::Unit::TestCase
             bucket.addfile("yayness", "/my/file")
         }
 
-        pathfile = File.join(bucket.path, sum, "paths")
+        a, b, pathfile = bucket.class.paths(bucket.path, sum)
 
         assert(FileTest.exists?(pathfile), "No path file at %s" % pathfile)
 
         assert_equal("/my/file\n", File.read(pathfile))
     end
+
+    # #447 -- a flat file structure just won't suffice.
+    def test_deeper_filestructure
+        bucket = Puppet::Network::Handler.filebucket.new(:Path => @bucket)
+
+        text = "this is some text"
+        md5 = Digest::MD5.hexdigest(text)
+
+        olddir = File.join(@bucket, md5)
+        FileUtils.mkdir_p(olddir)
+        oldcontent = File.join(olddir, "contents")
+        File.open(oldcontent, "w") { |f| f.print text }
+
+        result = nil
+        assert_nothing_raised("Could not retrieve content from old structure") do
+            result = bucket.getfile(md5)
+        end
+        assert_equal(text, result, "old-style content is wrong")
+
+        text = "and this is some new text"
+        md5 = Digest::MD5.hexdigest(text)
+
+        dirs = File.join(md5[0..7].split(""))
+        dir = File.join(@bucket, dirs)
+        filedir, contents, paths = bucket.class.paths(@bucket, md5)
+
+        assert_equal(dir, filedir, "did not use a deeper file structure") 
+        assert_equal(File.join(dir, "contents"), contents,
+            "content path is not the deeper version")
+        assert_equal(File.join(dir, "paths"), paths,
+            "paths file path is not the deeper version")
+
+        # Store our new text and make sure it gets stored in the new location
+        path = "/some/fake/path"
+        assert_nothing_raised("Could not store text") do
+            bucket.addfile(text, path)
+        end
+        assert(FileTest.exists?(contents), "did not create content file")
+        assert_equal(text, File.read(contents), "content is not right")
+        assert(FileTest.exists?(paths), "did not create paths file")
+        assert(File.read(paths).include?(path), "paths file does not contain path")
+
+        # And make sure we get it back out again
+        assert_nothing_raised("Could not retrieve new-style content") do
+            result = bucket.getfile(md5)
+        end
+        assert_equal(text, result, "did not retrieve new content correctly")
+    end
+
+    def test_add_path
+        bucket = Puppet::Network::Handler.filebucket.new(:Path => @bucket)
+
+        file = tempfile()
+
+        assert(! FileTest.exists?(file), "file already exists")
+
+        path = "/some/thing"
+        assert_nothing_raised("Could not add path") do
+            bucket.send(:add_path, path, file)
+        end
+        assert_equal(path + "\n", File.read(file), "path was not added")
+
+        assert_nothing_raised("Could not add path second time") do
+            bucket.send(:add_path, path, file)
+        end
+        assert_equal(path + "\n", File.read(file), "path was duplicated")
+
+        # Now try a new path
+        newpath = "/another/path"
+        assert_nothing_raised("Could not add path second time") do
+            bucket.send(:add_path, newpath, file)
+        end
+        text = [path, newpath].join("\n") + "\n"
+        assert_equal(text, File.read(file), "path was duplicated")
+
+        assert_nothing_raised("Could not add path third time") do
+            bucket.send(:add_path, path, file)
+        end
+        assert_equal(text, File.read(file), "path was duplicated")
+        assert_nothing_raised("Could not add second path second time") do
+            bucket.send(:add_path, newpath, file)
+        end
+        assert_equal(text, File.read(file), "path was duplicated")
+    end
 end
 
 # $Id$