Add top-level schema caches to Schema::Visibility for better performance

lib/graphql/schema.rb

@@ -446,7 +446,12 @@ def query(new_query_object = nil, &lazy_load_block)
             raise GraphQL::Error, "Second definition of `query(...)` (#{dup_defn.inspect}) is invalid, already configured with #{@query_object.inspect}"
           elsif use_visibility_profile?
             if block_given?
-              @query_object = lazy_load_block
+              if visibility.preload?
+                @query_object = lazy_load_block.call
+                self.visibility.query_configured(@query_object)
+              else
+                @query_object = lazy_load_block
+              end
             else
               @query_object = new_query_object
               self.visibility.query_configured(@query_object)
@@ -480,7 +485,12 @@ def mutation(new_mutation_object = nil, &lazy_load_block)
             raise GraphQL::Error, "Second definition of `mutation(...)` (#{dup_defn.inspect}) is invalid, already configured with #{@mutation_object.inspect}"
           elsif use_visibility_profile?
             if block_given?
-              @mutation_object = lazy_load_block
+              if visibility.preload?
+                @mutation_object = lazy_load_block.call
+                self.visibility.mutation_configured(@mutation_object)
+              else
+                @mutation_object = lazy_load_block
+              end
             else
               @mutation_object = new_mutation_object
               self.visibility.mutation_configured(@mutation_object)
@@ -514,7 +524,12 @@ def subscription(new_subscription_object = nil, &lazy_load_block)
             raise GraphQL::Error, "Second definition of `subscription(...)` (#{dup_defn.inspect}) is invalid, already configured with #{@subscription_object.inspect}"
           elsif use_visibility_profile?
             if block_given?
-              @subscription_object = lazy_load_block
+              if visibility.preload?
+                @subscription_object = lazy_load_block.call
+                visibility.subscription_configured(@subscription_object)
+              else
+                @subscription_object = lazy_load_block
+              end
             else
               @subscription_object = new_subscription_object
               self.visibility.subscription_configured(@subscription_object)

lib/graphql/schema/visibility.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 require "graphql/schema/visibility/profile"
 require "graphql/schema/visibility/migration"
+require "graphql/schema/visibility/visit"
 
 module GraphQL
   class Schema
@@ -13,40 +14,76 @@ class Visibility
       # @param migration_errors [Boolean] if `true`, raise an error when `Visibility` and `Warden` return different results
       def self.use(schema, dynamic: false, profiles: EmptyObjects::EMPTY_HASH, preload: (defined?(Rails) ? Rails.env.production? : nil), migration_errors: false)
         schema.visibility = self.new(schema, dynamic: dynamic, preload: preload, profiles: profiles, migration_errors: migration_errors)
+        if preload
+          schema.visibility.preload
+        end
       end
 
       def initialize(schema, dynamic:, preload:, profiles:, migration_errors:)
         @schema = schema
         schema.use_visibility_profile = true
-        if migration_errors
-          schema.visibility_profile_class = Migration
+        schema.visibility_profile_class = if migration_errors
+          Visibility::Migration
+        else
+          Visibility::Profile
         end
         @preload = preload
         @profiles = profiles
         @cached_profiles = {}
         @dynamic = dynamic
         @migration_errors = migration_errors
-        if preload
-          # Traverse the schema now (and in the *_configured hooks below)
-          # To make sure things are loaded during boot
-          @preloaded_types = Set.new
-          types_to_visit = [
-            @schema.query,
-            @schema.mutation,
-            @schema.subscription,
-            *@schema.introspection_system.types.values,
-            *@schema.introspection_system.entry_points.map { |ep| ep.type.unwrap },
-            *@schema.orphan_types,
-          ]
-          # Root types may have been nil:
-          types_to_visit.compact!
-          ensure_all_loaded(types_to_visit)
-
-          profiles.each do |profile_name, example_ctx|
-            example_ctx[:visibility_profile] = profile_name
-            prof = profile_for(example_ctx, profile_name)
-            prof.all_types # force loading
-          end
+        # Top-level type caches:
+        @visit = nil
+        @interface_type_memberships = nil
+        @directives = nil
+        @types = nil
+        @references = nil
+        @loaded_all = false
+      end
+
+      def all_directives
+        load_all
+        @directives
+      end
+
+      def all_interface_type_memberships
+        load_all
+        @interface_type_memberships
+      end
+
+      def all_references
+        load_all
+        @references
+      end
+
+      def get_type(type_name)
+        load_all
+        @types[type_name]
+      end
+
+      def preload?
+        @preload
+      end
+
+      def preload
+        # Traverse the schema now (and in the *_configured hooks below)
+        # To make sure things are loaded during boot
+        @preloaded_types = Set.new
+        types_to_visit = [
+          @schema.query,
+          @schema.mutation,
+          @schema.subscription,
+          *@schema.introspection_system.types.values,
+          *@schema.introspection_system.entry_points.map { |ep| ep.type.unwrap },
+          *@schema.orphan_types,
+        ]
+        # Root types may have been nil:
+        types_to_visit.compact!
+        ensure_all_loaded(types_to_visit)
+        @profiles.each do |profile_name, example_ctx|
+          example_ctx[:visibility_profile] = profile_name
+          prof = profile_for(example_ctx, profile_name)
+          prof.all_types # force loading
         end
       end
 
@@ -132,7 +169,10 @@ def profile_for(context, visibility_profile)
         end
       end
 
-      private
+      attr_reader :top_level
+
+      # @api private
+      attr_reader :unfiltered_interface_type_memberships
 
       def top_level_profile(refresh: false)
         if refresh
@@ -141,6 +181,8 @@ def top_level_profile(refresh: false)
         @top_level_profile ||= @schema.visibility_profile_class.new(context: Query::NullContext.instance, schema: @schema)
       end
 
+      private
+
       def ensure_all_loaded(types_to_visit)
         while (type = types_to_visit.shift)
           if type.kind.fields? && @preloaded_types.add?(type)
@@ -153,6 +195,86 @@ def ensure_all_loaded(types_to_visit)
         top_level_profile(refresh: true)
         nil
       end
+
+      def load_all(types: nil)
+        if @visit.nil?
+          # Set up the visit system
+          @interface_type_memberships = Hash.new { |h, interface_type| h[interface_type] = [] }.compare_by_identity
+          @directives = []
+          @types = {} # String => Module
+          @references = Hash.new { |h, member| h[member] = [] }.compare_by_identity
+          @unions_for_references = Set.new
+          @visit = Visibility::Visit.new(@schema) do |member|
+            if member.is_a?(Module)
+              type_name = member.graphql_name
+              if (prev_t = @types[type_name])
+                if prev_t.is_a?(Array)
+                  prev_t << member
+                else
+                  @types[type_name] = [member, prev_t]
+                end
+              else
+                @types[member.graphql_name] = member
+              end
+              if member < GraphQL::Schema::Directive
+                @directives << member
+              elsif member.respond_to?(:interface_type_memberships)
+                member.interface_type_memberships.each do |itm|
+                  @references[itm.abstract_type] << member
+                  @interface_type_memberships[itm.abstract_type] << itm
+                end
+              elsif member < GraphQL::Schema::Union
+                @unions_for_references << member
+              end
+            elsif member.is_a?(GraphQL::Schema::Argument)
+              member.validate_default_value
+              @references[member.type.unwrap] << member
+            elsif member.is_a?(GraphQL::Schema::Field)
+              @references[member.type.unwrap] << member
+            end
+            true
+          end
+
+          @schema.root_types.each do |t|
+            @references[t] << true
+          end
+
+          @schema.introspection_system.types.each_value do |t|
+            @references[t] << true
+          end
+          @visit.visit_each(types: []) # visit default directives
+        end
+
+        if types
+          @visit.visit_each(types: types, directives: [])
+        elsif @loaded_all == false
+          @loaded_all = true
+          @visit.visit_each
+        else
+          # already loaded all
+          return
+        end
+
+        # TODO: somehow don't iterate over all these,
+        # only the ones that may have been modified
+        @interface_type_memberships.each do |int_type, type_memberships|
+          referers = @references[int_type].select { |r| r.is_a?(GraphQL::Schema::Field) }
+          if referers.any?
+            type_memberships.each do |type_membership|
+              implementor_type = type_membership.object_type
+              # Add new items only:
+              @references[implementor_type] |= referers
+            end
+          end
+        end
+
+        @unions_for_references.each do |union_type|
+          refs = @references[union_type]
+          union_type.all_possible_types.each do |object_type|
+            @references[object_type] |= refs # Add new items
+          end
+        end
+      end
     end
   end
 end