Add a batch loading benchmark

Gemfile

@@ -7,6 +7,7 @@ gem 'bootsnap' # required by the Rails apps generated in tests
 gem 'ruby-prof', platform: :ruby
 gem 'pry'
 gem 'pry-stack_explorer', platform: :ruby
+gem 'graphql-batch'
 if RUBY_VERSION >= "2.4"
   gem 'pry-byebug'
 end

Rakefile

@@ -98,6 +98,12 @@ namespace :bench do
     prepare_benchmark
     GraphQLBenchmark.profile_large_result
   end
+
+  desc "Compare GraphQL-Batch and GraphQL-Dataloader"
+  task :profile_batch_loaders do
+    prepare_benchmark
+    GraphQLBenchmark.profile_batch_loaders
+  end
 end
 
 namespace :test do

benchmark/batch_loading.rb

@@ -0,0 +1,144 @@
+module BatchLoading
+  class GraphQLBatchSchema < GraphQL::Schema
+    DATA = [
+      { id: "1", name: "Bulls", player_ids: ["2", "3"] },
+      { id: "2", name: "Michael Jordan", team_id: "1" },
+      { id: "3", name: "Scottie Pippin", team_id: "1" },
+      { id: "4", name: "Braves", player_ids: ["5", "6"] },
+      { id: "5", name: "Chipper Jones", team_id: "4" },
+      { id: "6", name: "Tom Glavine", team_id: "4" },
+    ]
+
+    class DataLoader < GraphQL::Batch::Loader
+      def initialize(column: :id)
+        @column = column
+      end
+
+      def perform(keys)
+        keys.each do |key|
+          record = DATA.find { |d| d[@column] == key }
+          fulfill(key, record)
+        end
+      end
+    end
+
+    class Team < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :players, "[BatchLoading::GraphQLBatchSchema::Player]", null: false
+
+      def players
+        DataLoader.load_many(object[:player_ids])
+      end
+    end
+
+    class Player < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :team, Team, null: false
+
+      def team
+        DataLoader.load(object[:team_id])
+      end
+    end
+
+    class Query < GraphQL::Schema::Object
+      field :team, Team, null: true do
+        argument :name, String, required: true
+      end
+
+      def team(name:)
+        DataLoader.for(column: :name).load(name)
+      end
+    end
+
+    query(Query)
+    use GraphQL::Execution::Interpreter
+    use GraphQL::Analysis::AST
+    use GraphQL::Batch
+  end
+
+  class GraphQLDataloaderSchema < GraphQL::Schema
+    class DataSource < GraphQL::Dataloader::Source
+      def initialize(dataloader, column: :id)
+        @column = column
+        super(dataloader)
+      end
+
+      def fetch(keys)
+        keys.map { |key|
+          d = GraphQLBatchSchema::DATA.find { |d| d[@column] == key }
+          # p [key, @column, d]
+          d
+        }
+      end
+    end
+
+    class Team < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :players, "[BatchLoading::GraphQLDataloaderSchema::Player]", null: false
+
+      def players
+        dataloader.with(DataSource).load_all(object[:player_ids])
+      end
+    end
+
+    class Player < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :team, Team, null: false
+
+      def team
+        dataloader.with(DataSource).load(object[:team_id])
+      end
+    end
+
+    class Query < GraphQL::Schema::Object
+      field :team, Team, null: true do
+        argument :name, String, required: true
+      end
+
+      def team(name:)
+        dataloader.with(DataSource, column: :name).load(name)
+      end
+    end
+
+    query(Query)
+    use GraphQL::Execution::Interpreter
+    use GraphQL::Analysis::AST
+    use GraphQL::Dataloader
+  end
+
+  class GraphQLNoBatchingSchema < GraphQL::Schema
+    DATA = GraphQLBatchSchema::DATA
+
+    class Team < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :players, "[BatchLoading::GraphQLNoBatchingSchema::Player]", null: false
+
+      def players
+        object[:player_ids].map { |id| DATA.find { |d| d[:id] == id } }
+      end
+    end
+
+    class Player < GraphQL::Schema::Object
+      field :name, String, null: false
+      field :team, Team, null: false
+
+      def team
+        DATA.find { |d| d[:id] == object[:team_id] }
+      end
+    end
+
+    class Query < GraphQL::Schema::Object
+      field :team, Team, null: true do
+        argument :name, String, required: true
+      end
+
+      def team(name:)
+        DATA.find { |d| d[:name] == name }
+      end
+    end
+
+    query(Query)
+    use GraphQL::Execution::Interpreter
+    use GraphQL::Analysis::AST
+  end
+end

benchmark/run.rb

@@ -5,6 +5,7 @@
 require "benchmark/ips"
 require "ruby-prof"
 require "memory_profiler"
+require "graphql/batch"
 
 module GraphQLBenchmark
   QUERY_STRING = GraphQL::Introspection::INTROSPECTION_QUERY
@@ -143,4 +144,50 @@ class Schema < GraphQL::Schema
       }
     GRAPHQL
   end
+
+  def self.profile_batch_loaders
+    require_relative "./batch_loading"
+    include BatchLoading
+
+    document = GraphQL.parse <<-GRAPHQL
+    {
+      braves: team(name: "Braves") { ...TeamFields }
+      bulls: team(name: "Bulls") { ...TeamFields }
+    }
+
+    fragment TeamFields on Team {
+      players {
+        team {
+          players {
+            team {
+              name
+            }
+          }
+        }
+      }
+    }
+    GRAPHQL
+    batch_result = GraphQLBatchSchema.execute(document: document).to_h
+    dataloader_result = GraphQLDataloaderSchema.execute(document: document).to_h
+    no_batch_result = GraphQLNoBatchingSchema.execute(document: document).to_h
+
+    results = [batch_result, dataloader_result, no_batch_result].uniq
+    if results.size > 1
+      puts "Batch result:"
+      pp batch_result
+      puts "Dataloader result:"
+      pp dataloader_result
+      puts "No-batch result:"
+      pp no_batch_result
+      raise "Got different results -- fix implementation before benchmarking."
+    end
+
+    Benchmark.ips do |x|
+      x.report("GraphQL::Batch") { GraphQLBatchSchema.execute(document: document) }
+      x.report("GraphQL::Dataloader") { GraphQLDataloaderSchema.execute(document: document) }
+      x.report("No Batching") { GraphQLNoBatchingSchema.execute(document: document) }
+
+      x.compare!
+    end
+  end
 end

lib/graphql/dataloader.rb

@@ -31,7 +31,7 @@ def self.use(schema)
 
     def initialize(context)
       @context = context
-      @source_cache = {}
+      @source_cache = Hash.new { |h,k| h[k] = {} }
       @waiting_fibers = []
     end
 
@@ -117,8 +117,8 @@ def run
       nil
     end
 
-    def with(source_class)
-      @source_cache[source_class] ||= source_class.new(self)
+    def with(source_class, *batch_parameters)
+      @source_cache[source_class][batch_parameters] ||= source_class.new(self, *batch_parameters)
     end
 
     private
@@ -129,8 +129,17 @@ def with(source_class)
     def create_source_fiber_stack
       # only assign a new array when we need one:
       source_fiber_stack = nil
-      while @source_cache.each_value.any?(&:pending?)
-        pending_sources = @source_cache.each_value.select(&:pending?)
+      pending_sources = nil
+      @source_cache.each_value do |source_by_batch_params|
+        source_by_batch_params.each_value do |source|
+          if source.pending?
+            pending_sources ||= []
+            pending_sources << source
+          end
+        end
+      end
+
+      if pending_sources
         source_fiber = Fiber.new do
           pending_sources.each(&:run_pending_keys)
         end