Merge pull request #1 from acidprime/master

Update from original

Author: vicinus

README.md

@@ -98,6 +98,12 @@ to the `--threads` option. This will balence the catalogs evenly on the old and
 masters. This option defaults to 10 and in testing 50 threads seemed correct for
 4 masters with two load balancers.
 
+Note: When using catalog diff to compare directories, one thread per catalog
+comparison will be created.  However, since Ruby cannot take advantage of
+multiple CPUs this may be of limited use comparing local catalogs.  If the
+'parallel' gem is installed, then one process will be forked off per CPU on the
+system, allowing use of all CPUs.
+
 ## Fact search
 You can pass `--fact_search` to filter the list of nodes based on a single fact value.
 This currently defaults to `kernel=Linux` if you do not pass it. The yaml cache will be

lib/puppet/catalog-diff/differ.rb

@@ -40,6 +40,12 @@ def diff(options = {})
           tmp = Marshal.load(File.read(r))
         when '.pson'
           tmp = PSON.load(File.read(r))
+          unless tmp.respond_to? :version
+            tmp = PSON.load({
+              'document_type' => 'Catalog',
+              'data' => PSON.load(File.read(r)),
+            }.to_pson)
+          end
         when '.json'
           tmp = PSON.load(File.read(r))
 	else

lib/puppet/catalog-diff/formater.rb

@@ -5,27 +5,57 @@ class Formater
     def initialize()
     end
 
+    def format_simple(v, indent='', do_indent=false, comma='')
+      str = ''
+      str << indent if do_indent
+      v = "\"#{v}\"" unless [Fixnum, TrueClass, FalseClass].include?(v.class)
+      str << v.to_s << comma << "\n"
+    end
+
+    def format_array(v, indent='', do_indent=false, comma='')
+      str = ''
+      str << indent if do_indent
+      str << '[' << "\n"
+      v.each do |val|
+        str << format_value(val, "#{indent}     ", true, ',')
+      end
+      str << "\t     #{indent}    ]" << "\n" << comma
+    end
+
+    def format_hash(v, indent='', do_indent=false, comma='')
+      str = ''
+      str << indent if do_indent
+      str << '{' << "\n"
+      v.each do |key, val|
+        str << "\t     #{indent}     #{key} => "
+        str << format_value(val, "#{indent}     ", true, ',', key)
+      end
+      str << "\t     #{indent}    }" << "\n" << comma
+    end
+
+    def format_value(v, indent='', do_indent=false, comma='', k=nil)
+      if v.is_a?(Array)
+        format_array(v, indent, do_indent, comma)
+      elsif v.is_a?(Hash)
+        format_hash(v, indent, do_indent, comma)
+      else
+        v = v[:checksum] if (k == :content && v.is_a?(Hash))
+        format_simple(v, indent, do_indent, comma)
+      end
+    end
+
     # creates a string representation of a resource that looks like Puppet code
     def resource_to_string(resource)
       str = ''
       str << "\t" + resource[:type].downcase << '{"' <<  resource[:title].to_s << '":' << "\n"
       params = Hash[(resource[:parameters].sort_by {|k, v| k})]
       params.each_pair do |k,v|
-        if v.is_a?(Array)
-          indent = " " * k.to_s.size
-          str << "\t     #{k} => [" << "\n"
-          v.each do |val|
-            str << "\t     #{indent}     #{val}," << "\n"
-          end
-          str << "\t     #{indent}    ]" << "\n"
-        else
-          if k == :content
-            v = v[:checksum]
-          end
-          str << "\t     #{k} => #{v}" << "\n"
-        end
+        str << "\t     #{k} => "
+        indent = " " * k.to_s.size
+        str << format_value(v, indent, false, '', k)
       end
       str << "\t}\n"
+      str
     end
 
     def node_summary_header(node,summary,key)

lib/puppet/catalog-diff/searchfacts.rb

@@ -73,7 +73,24 @@ def find_nodes_puppetdb(env)
         connection = Puppet::Network::HttpPool.http_instance(Puppet::Util::Puppetdb.server,port,use_ssl)
         base_query = ["and", ["=", ["node","active"], true]]
         base_query.concat([["=", "catalog-environment", env]]) if env
-        query = base_query.concat(@facts.map { |k, v| ["=", ["fact", k], v] })
+        real_facts = @facts.select { |k, v| !v.nil? }
+        query = base_query.concat(real_facts.map { |k, v| ["=", ["fact", k], v] })
+        classes = Hash[@facts.select { |k, v| v.nil? }].keys
+        classes.each do |c|
+          capit = c.split('::').map{ |n| n.capitalize }.join('::')
+          query = query.concat(
+            [["in", "certname",
+              ["extract", "certname",
+                ["select-resources",
+                  ["and",
+                    ["=", "type", "Class"],
+                    ["=", "title", capit ],
+                  ],
+                ],
+              ],
+            ]]
+          )
+        end
         json_query = URI.escape(query.to_json)
         unless filtered = PSON.load(connection.request_get("/v4/nodes/?query=#{json_query}", {"Accept" => 'application/json'}).body)
           raise "Error parsing json output of puppet search"

lib/puppet/face/catalog/diff.rb

@@ -1,6 +1,14 @@
 require 'puppet/face'
 require  'thread'
 require 'json'
+
+begin
+  require 'parallel'
+  HAS_PARALLEL_GEM = true
+rescue LoadError
+  HAS_PARALLEL_GEM = false
+end
+
 Puppet::Face.define(:catalog, '0.0.1') do
 
   action :diff do
@@ -103,19 +111,29 @@
         found_catalogs = Puppet::CatalogDiff::FindCatalogs.new(catalog1,catalog2).return_catalogs(options)
         new_catalogs = found_catalogs.keys
 
-        thread_count = 1
-        mutex = Mutex.new
-
-        thread_count.times.map {
-          Thread.new(nodes,new_catalogs,options) do |nodes,new_catalogs,options|
-            while new_catalog = mutex.synchronize { new_catalogs.pop }
-              node_name    = File.basename(new_catalog,File.extname(new_catalog))
-              old_catalog  = found_catalogs[new_catalog]
-              node_summary = Puppet::CatalogDiff::Differ.new(old_catalog, new_catalog).diff(options)
-              mutex.synchronize { nodes[node_name] = node_summary }
-            end
+        if HAS_PARALLEL_GEM
+          results = Parallel.map(new_catalogs) do |new_catalog|
+            node_name    = File.basename(new_catalog,File.extname(new_catalog))
+            old_catalog  = found_catalogs[new_catalog]
+            node_summary = Puppet::CatalogDiff::Differ.new(old_catalog, new_catalog).diff(options)
+            [ node_name, node_summary ]
           end
-        }.each(&:join)
+          nodes = Hash[results]
+        else
+          thread_count = 1
+          mutex = Mutex.new
+
+          thread_count.times.map {
+            Thread.new(nodes,new_catalogs,options) do |nodes,new_catalogs,options|
+              while new_catalog = mutex.synchronize { new_catalogs.pop }
+                node_name    = File.basename(new_catalog,File.extname(new_catalog))
+                old_catalog  = found_catalogs[new_catalog]
+                node_summary = Puppet::CatalogDiff::Differ.new(old_catalog, new_catalog).diff(options)
+                mutex.synchronize { nodes[node_name] = node_summary }
+              end
+            end
+          }.each(&:join)
+        end
       elsif File.file?(catalog1) && File.file?(catalog2)
         # User passed us two files
         node_name = File.basename(catalog2,File.extname(catalog2))
@@ -128,6 +146,8 @@
         pull_output = Puppet::Face[:catalog, '0.0.1'].pull(old_catalogs,new_catalogs,options[:fact_search],:old_server => catalog1,:new_server => catalog2,:changed_depth => options[:changed_depth], :threads => options[:threads], :use_puppetdb => options[:use_puppetdb])
         diff_output = Puppet::Face[:catalog, '0.0.1'].diff(old_catalogs,new_catalogs,options)
         nodes = diff_output
+        FileUtils.rm_rf(old_catalogs)
+        FileUtils.rm_rf(new_catalogs)
         nodes[:pull_output] = pull_output
         # Save the file as it can take a while to create
         if options[:output_report]