Use MiqPreloader/Relationship improvements in MiqRequestWorkflow

app/models/miq_request_workflow.rb

@@ -942,7 +942,15 @@ def get_ems_metadata_tree(src)
     @ems_metadata_tree ||= begin
       return if src[:ems].nil?
       st = Time.zone.now
-      result = load_ar_obj(src[:ems]).fulltree_arranged(:except_type => "VmOrTemplate")
+      hosts_scope   = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+      fulltree_opts = {
+        :except_type               => "VmOrTemplate",
+        :class_specific_preloaders => {
+          EmsCluster => [:hosts, hosts_scope],
+          Host       => hosts_scope
+        }
+      }
+      result = load_ar_obj(src[:ems]).fulltree_arranged(fulltree_opts)
       ems_metadata_tree_add_hosts_under_clusters!(result)
       @ems_xml_nodes = {}
       xml = MiqXml.newDoc(:xmlhash)

app/models/mixins/relationship_mixin.rb

@@ -515,7 +515,8 @@ def fulltree_ids_arranged(*args)
 
   # Returns the records in the tree from the root arranged in a tree
   def fulltree_arranged(*args)
-    Relationship.arranged_rels_to_resources(fulltree_rels_arranged(*args))
+    class_specific_preloaders = args.last.delete(:class_specific_preloaders) if args.last.kind_of?(Hash)
+    Relationship.arranged_rels_to_resources(fulltree_rels_arranged(*args), true, class_specific_preloaders)
   end
 
   # Returns a list of all unique child types

app/models/relationship.rb

@@ -82,8 +82,11 @@ def self.flatten_arranged_rels(relationships)
     end
   end
 
-  def self.arranged_rels_to_resources(relationships, initial = true)
-    MiqPreloader.preload(flatten_arranged_rels(relationships), :resource) if initial
+  def self.arranged_rels_to_resources(relationships, initial = true, class_specific_preloaders = nil)
+    if initial
+      record_set = flatten_arranged_rels(relationships)
+      MiqPreloader.polymorphic_preload_for_child_classes(record_set, :resource, class_specific_preloaders)
+    end
 
     relationships.each_with_object({}) do |(rel, children), h|
       h[rel.resource] = arranged_rels_to_resources(children, false)

lib/miq_preloader.rb

@@ -31,4 +31,88 @@ def self.preload_and_scope(records, association_name)
       target_klass.where(join_key.key.to_sym => records.select(join_key.foreign_key.to_sym))
     end
   end
+
+  # Allows having a polymorphic preloader, but then having class specific
+  # preloaders fire for the loaded polymorphic classes.
+  #
+  # @param records [ActiveRecord::Relation] collection of activerecord objects to preload into
+  # @param associations [Symbol|String|Array|Hash] association(s) to load (see .includes for examples)
+  # @param class_preloaders [Hash] keys are Classes, and values are associations for that polymorphic type
+  #
+  # Values for class_preloaders can either be an Array of two (args for sub
+  # preloader relationships), or an arel statement, which is the arel scope to
+  # execute for that specific class only (no sub relationships preloaded.
+  #
+  # Example:
+  #
+  #   irb> tree = ExtManagementSystem.last.fulltree_rels_arranged(:except_type => "VmOrTemplate")
+  #   irb> records = Relationship.flatten_arranged_rels(tree)
+  #   irb> hosts_scope = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+  #   irb> preloaders_per_class = { EmsCluster => [:hosts, hosts_scope], Host => hosts_scope }
+  #   irb> MiqPreloader.polymorphic_preload_for_child_classes(records, nil, preloaders_per_class)
+  #
+  # Note:  Class's .base_class are favored over their specific class
+  #
+  def self.polymorphic_preload_for_child_classes(records, associations, class_preloaders = {})
+    preloader = ActiveRecord::Associations::Preloader.new
+    preloader.extend(Module.new {
+      attr_accessor :class_specific_preloaders
+
+      # DIRTY HACK... but hey... at least I am just isolating it to this method
+      # right...
+      #
+      # Everyone else: "No Nick... just no..."
+
+      # Updated form from ar_virtual.rb, and merged with the code originally in
+      # ActiveRecord.  If the code in ar_virtual.rb is changed, this should
+      # probably also be updated.
+      def preloaders_for_one(association, records, scope)
+        klass_map = records.compact.group_by(&:class)
+
+        loaders = klass_map.keys.group_by { |klass| klass.virtual_includes(association) }.flat_map do |virtuals, klasses|
+          subset = klasses.flat_map { |klass| klass_map[klass] }
+          preload(subset, virtuals)
+        end
+
+        records_with_association = klass_map.select { |k, rs| k.reflect_on_association(association) }.flat_map { |k, rs| rs }
+        if records_with_association.any?
+          # This injects the original code from preloaders_for_one from
+          # ActiveRecord so we can add our own `if` in the middle of it.  The
+          # positive part of the `if` is the only portion of this that has
+          # changed, and the code copied is within the `loaders.concat`.
+          loaders.concat(grouped_records(association, records_with_association).flat_map do |reflection, klasses|
+            klasses.map do |rhs_klass, rs|
+              base_klass = rhs_klass.base_class if rhs_klass.respond_to?(:base_class)
+
+              # Start of new code (1)
+              class_preloader = (class_specific_preloaders[base_klass] || class_specific_preloaders[rhs_klass])
+              loader_klass    = preloader_for(reflection, rs, rhs_klass)
+
+              loader = if class_preloader.kind_of?(ActiveRecord::Relation)
+                         loader_klass.new(rhs_klass, rs, reflection, class_preloader)
+                       else
+                         loader_klass.new(rhs_klass, rs, reflection, scope)
+                       end
+              # End of new code (1)
+
+              loader.run self
+
+              # Start of new code (2)
+              if class_preloader.kind_of?(Array) && class_preloader.count == 2
+                [loader, MiqPreloader.preload(loader.preloaded_records, class_preloader[0], class_preloader[1])]
+              else
+                loader
+              end
+              # End of new code (2)
+            end
+          end)
+        end
+
+        loaders
+      end
+    })
+    preloader.class_specific_preloaders = class_preloaders || {}
+
+    preloader.preload(records, associations, nil)
+  end
 end

spec/lib/miq_preloader_spec.rb

@@ -112,4 +112,46 @@ def preload_and_scope(*args)
       MiqPreloader.preload_and_scope(*args)
     end
   end
+
+  describe ".polymorphic_preload_for_child_classes" do
+    include_context "simple ems_metadata tree" do
+      before { init_full_tree }
+    end
+
+    it "preloads polymorphic relationships that are defined" do
+      tree    = ems.fulltree_rels_arranged(:except_type => "VmOrTemplate")
+      records = Relationship.flatten_arranged_rels(tree)
+
+      hosts_scope   = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+      class_loaders = { EmsCluster => [:hosts, hosts_scope], Host => hosts_scope }
+
+      # 4 queries are expected here:
+      #
+      # - 1 for ExtManagementSystem (root)
+      # - 1 for EmsClusters in the tree
+      # - 1 for Hosts in the tree
+      # - 1 for Hosts from relation in EmsClusters
+      #
+      # Since all the hosts in this case are also part of the tree, there are
+      # "duplicate hosts loaded", but that was the nature of this prior to the
+      # change anyway, so this is not new.  This does make it soe that any
+      # hosts accessed through a EmsCluster are preloaded, however, instead of
+      # an N+1.
+      #
+      # In some cases, a ems_metatdata tree might not include all of the
+      # hosts of a EMS, but some still exist as part of the cluster.  This also
+      # makes sure both cases are covered, and the minimal amount of queries
+      # are still executed.
+      #
+      # rubocop:disable Style/BlockDelimiters
+      expect {
+        MiqPreloader.polymorphic_preload_for_child_classes(records, :resource, class_loaders)
+        records.select   { |rel|  rel.resource if rel.resource_type == "Host" }
+               .each     { |rel|  rel.resource.v_total_vms }
+        records.select   { |rel|  rel.resource if rel.resource_type == "EmsCluster" }
+               .flat_map { |rel|  rel.resource.hosts }.each(&:v_total_vms)
+      }.to match_query_limit_of(4)
+      # rubocop:enable Style/BlockDelimiters
+    end
+  end
 end

spec/models/mixins/relationship_mixin_spec.rb

@@ -747,6 +747,64 @@
         }
       )
     end
+
+    context "with a EMS based tree" do
+      # Note:  This shared_context overwrites the `let(:vms)` at the top of
+      # this file.
+      include_context "simple ems_metadata tree" do
+        before { init_full_tree }
+      end
+
+      it "builds the tree normally" do
+        nodes = ems.fulltree_arranged
+        expect(nodes).to eq(
+          ems => {
+            clusters[0] => {
+              hosts[0] => {
+                vms[0] => {},
+                vms[1] => {}
+              },
+              hosts[1] => {
+                vms[2] => {},
+                vms[3] => {}
+              }
+            },
+            clusters[1] => {
+              hosts[2] => {
+                vms[4] => {},
+                vms[5] => {}
+              },
+              hosts[3] => {
+                vms[6] => {},
+                vms[7] => {}
+              }
+            }
+          }
+        )
+      end
+
+      it "can preload certain relationships to avoid N+1s" do
+        hosts_scope   = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+        fulltree_opts = {
+          :except_type               => "VmOrTemplate",
+          :class_specific_preloaders => {
+            EmsCluster => [:hosts, hosts_scope],
+            Host       => hosts_scope
+          }
+        }
+        tree = ems.fulltree_arranged(fulltree_opts)
+
+        # rubocop:disable Style/BlockDelimiters
+        expect {
+          # get the v_total_vms through the EmsCluster records
+          tree.values.first.keys.flat_map(&:hosts).each(&:v_total_vms)
+
+          # get the v_total_vms through the Host records in the tree
+          tree.values.first.values.flat_map(&:keys).each(&:v_total_vms)
+        }.to match_query_limit_of(0)
+        # rubocop:enable Style/BlockDelimiters
+      end
+    end
   end
 
   describe "#fulltree_ids_arranged" do

spec/support/examples_group/shared_context_for_ems_metadata_tree.rb

@@ -0,0 +1,68 @@
+shared_context "simple ems_metadata tree" do
+  #####################
+  # Context Variables #
+  #####################
+
+  # These can be updated to increase the amount of data in the tree
+
+  let(:cluster_count)     { 2 }
+  let(:hosts_per_cluster) { 2 }
+  let(:vms_per_host)      { 2 }
+
+  ###############
+  # Base models #
+  ###############
+
+  let(:ems)      { FactoryGirl.create(:ems_infra) }
+  let(:clusters) { FactoryGirl.create_list(:ems_cluster, cluster_count, :ext_management_system => ems) }
+
+  let(:hosts) do
+    hosts = []
+    clusters.each do |cluster|
+      hosts += FactoryGirl.create_list(:host, hosts_per_cluster,
+                                       :ext_management_system => ems,
+                                       :ems_cluster           => cluster)
+    end
+    hosts
+  end
+
+  let(:vms) do
+    vms = []
+    hosts.each do |host|
+      vms += FactoryGirl.create_list(:vm, vms_per_host,
+                                     :ext_management_system => ems,
+                                     :host                  => host)
+    end
+    vms
+  end
+
+  #################
+  # Relationships #
+  #################
+
+  let(:ems_rel)      { ems.init_relationship }
+  let(:cluster_rels) { clusters.map { |cluster| cluster.init_relationship(ems_rel) } }
+
+  # The next to use a integer division trick to map the proper parent_rel to
+  # the record being created (the `[index / child_per_parent]` part)
+
+  let(:host_rels) do
+    hosts.map.with_index do |host, index|
+      host.init_relationship(cluster_rels[index / hosts_per_cluster])
+    end
+  end
+
+  let(:vm_rels) do
+    vms.map.with_index do |vm, index|
+      vm.init_relationship(host_rels[index / vms_per_host])
+    end
+  end
+
+  ###########
+  # Helpers #
+  ###########
+
+  # Convenience statement for initializing the tree (there is no included
+  # `before` to do this automatically
+  let(:init_full_tree) { vm_rels }
+end