Improve GraphQL::Dataloader::Source#sync's efficiency

[glasses618]

Summary

In the discussion #4512, we discovered an issue with sync related to

graphql-ruby/lib/graphql/dataloader/source.rb

Line 74 in 250a724

while pending_keys.any? { |k| !@results.key?(k) }

When the pending_keys, which is a duplicate of @pending_keys, becomes large, it tends to consume much Ruby CPU resources, particularly when dealing with ActiveRecord as a Source key.

This PR modify the function of sync so that it only waits for the key on which load was called.

graphql-ruby/lib/graphql/dataloader/source.rb

Line 37 in 250a724

def load(key)

graphql-ruby/lib/graphql/dataloader/source.rb

Line 49 in 250a724

def load_all(keys)

Benchmark

In the benchmark, I use Struct as a Source key , and load 5000 Struct objects.
In Ruby 2.7.8, this approach is ~9.49x faster than before but it increases ~0.21% memory. In Ruby 3.0.6, this approach is ~4.99x faster than before but it increases ~0.22% memory.

Benchmark Script

require "graphql"
require "benchmark/ips"
require "graphql/batch"
require "memory_profiler"

USER_STRUCT = Struct.new(:id, keyword_init: true)
POST_STRUCT = Struct.new(:id, :user_id, keyword_init: true)

class Data
  USERS_DATA = (1..5000).map{|s| USER_STRUCT.new(id: s) }.freeze
  POSTS_DATA = (1..5000).map{|s| POST_STRUCT.new(id: s, user_id: s) }.freeze
end

class GraphQLBatchSchema < GraphQL::Schema
  class DataLoader < GraphQL::Batch::Loader
    def perform(posts)
      posts.each do |post|
        user = Data::USERS_DATA.find { |u| u.id == post.user_id }
        fulfill(post, user)
      end
    end
  end

  class User < GraphQL::Schema::Object
    field :id, String, null: false
  end

  class Post < GraphQL::Schema::Object
    field :user, User, null: false

    def user
      DataLoader.load(object)
    end
  end

  class Query < GraphQL::Schema::Object
    field :posts, [Post]

    def posts
      Data::POSTS_DATA
    end
  end

  query(Query)
  use GraphQL::Batch
end

class GraphQLDataloaderSchema < GraphQL::Schema
  class DataSource < GraphQL::Dataloader::Source
    def fetch(posts)
      posts.map do |post|
        Data::USERS_DATA.find { |u| u.id == post.user_id }
      end
    end
  end

  class User < GraphQL::Schema::Object
    field :id, String, null: false
  end

  class Post < GraphQL::Schema::Object
    field :user, User, null: false

    def user
      dataloader.with(DataSource).load(object)
    end
  end

  class Query < GraphQL::Schema::Object
    field :posts, [Post]

    def posts
      Data::POSTS_DATA
    end
  end

  query(Query)
  use GraphQL::Dataloader
end

class GraphQLNoBatchingSchema < GraphQL::Schema
  class User < GraphQL::Schema::Object
    field :id, String, null: false
  end

  class Post < GraphQL::Schema::Object
    field :user, User, null: false

    def user
      Data::USERS_DATA.find { |u| u.id == object.user_id }
    end
  end

  class Query < GraphQL::Schema::Object
    field :posts, [Post]

    def posts
      Data::POSTS_DATA
    end
  end
  query(Query)
end
document = GraphQL.parse <<-GRAPHQL
{
  posts {
    user {
      id
    }
  }
}
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

puts "========== GraphQL-Batch Memory =============="
report = MemoryProfiler.report do
  GraphQLBatchSchema.execute(document: document)
end

report.pretty_print

puts "========== Dataloader Memory ================="
report = MemoryProfiler.report do
  GraphQLDataloaderSchema.execute(document: document)
end

report.pretty_print

puts "========== No Batch Memory =============="
report = MemoryProfiler.report do
  GraphQLNoBatchingSchema.execute(document: document)
end

report.pretty_print

Ruby 2.7.8

Before

Warming up --------------------------------------
      GraphQL::Batch     1.000  i/100ms
 GraphQL::Dataloader     1.000  i/100ms
         No Batching     1.000  i/100ms
Calculating -------------------------------------
      GraphQL::Batch      1.490  (± 0.0%) i/s -      8.000  in   5.368523s
 GraphQL::Dataloader      0.154  (± 0.0%) i/s -      1.000  in   6.497724s
         No Batching      1.423  (± 0.0%) i/s -      8.000  in   5.623860s

Comparison:
      GraphQL::Batch:        1.5 i/s
         No Batching:        1.4 i/s - 1.05x  slower
 GraphQL::Dataloader:        0.2 i/s - 9.68x  slower

========== GraphQL-Batch Memory ==============
Total allocated: 9081112 bytes (95298 objects)
Total retained:  0 bytes (0 objects)
...
========== Dataloader Memory =================
Total allocated: 27287560 bytes (175328 objects)
Total retained:  168 bytes (1 objects)
...
========== No Batch Memory ==============
Total allocated: 5731328 bytes (50290 objects)
Total retained:  0 bytes (0 objects)
...

After

Warming up --------------------------------------
      GraphQL::Batch     1.000  i/100ms
 GraphQL::Dataloader     1.000  i/100ms
         No Batching     1.000  i/100ms
Calculating -------------------------------------
      GraphQL::Batch      1.495  (± 0.0%) i/s -      8.000  in   5.352489s
 GraphQL::Dataloader      1.462  (± 0.0%) i/s -      8.000  in   5.470723s
         No Batching      1.420  (± 0.0%) i/s -      8.000  in   5.632891s

Comparison:
      GraphQL::Batch:        1.5 i/s
 GraphQL::Dataloader:        1.5 i/s - 1.02x  slower
         No Batching:        1.4 i/s - 1.05x  slower

========== GraphQL-Batch Memory ==============
Total allocated: 9081112 bytes (95298 objects)
Total retained:  0 bytes (0 objects)
...
========== Dataloader Memory =================
Total allocated: 27347336 bytes (175328 objects)
Total retained:  168 bytes (1 objects)
...
========== No Batch Memory ==============
Total allocated: 5731328 bytes (50290 objects)
Total retained:  0 bytes (0 objects)
...

Ruby 3.0.6

Before

Warming up --------------------------------------
      GraphQL::Batch     1.000  i/100ms
 GraphQL::Dataloader     1.000  i/100ms
         No Batching     1.000  i/100ms
Calculating -------------------------------------
      GraphQL::Batch      1.408  (± 0.0%) i/s -      8.000  in   5.682068s
 GraphQL::Dataloader      0.275  (± 0.0%) i/s -      2.000  in   7.264746s
         No Batching      1.292  (± 0.0%) i/s -      7.000  in   5.418212s

Comparison:
      GraphQL::Batch:        1.4 i/s
         No Batching:        1.3 i/s - 1.09x  slower
 GraphQL::Dataloader:        0.3 i/s - 5.11x  slower

========== GraphQL-Batch Memory ==============
Total allocated: 9080608 bytes (95295 objects)
Total retained:  0 bytes (0 objects)
...
========== Dataloader Memory =================
Total allocated: 27046472 bytes (195327 objects)
Total retained:  0 bytes (0 objects)
...
========== No Batch Memory ==============
Total allocated: 5730824 bytes (50287 objects)
Total retained:  0 bytes (0 objects)
...

After

Warming up --------------------------------------
      GraphQL::Batch     1.000  i/100ms
 GraphQL::Dataloader     1.000  i/100ms
         No Batching     1.000  i/100ms
Calculating -------------------------------------
      GraphQL::Batch      1.408  (± 0.0%) i/s -      8.000  in   5.682212s
 GraphQL::Dataloader      1.371  (± 0.0%) i/s -      7.000  in   5.106177s
         No Batching      1.296  (± 0.0%) i/s -      7.000  in   5.400772s

Comparison:
      GraphQL::Batch:        1.4 i/s
 GraphQL::Dataloader:        1.4 i/s - 1.03x  slower
         No Batching:        1.3 i/s - 1.09x  slower

========== GraphQL-Batch Memory ==============
Total allocated: 9080608 bytes (95295 objects)
Total retained:  3635680 bytes (30085 objects)
...
========== Dataloader Memory =================
Total allocated: 27106248 bytes (195327 objects)
Total retained:  0 bytes (0 objects)
...
========== No Batch Memory ==============
Total allocated: 5730824 bytes (50287 objects)
Total retained:  0 bytes (0 objects)
...

[glasses618]

Sorry, I found I don't need to dup the pending_keys which sync was called. Let me make some changes. https://github.com/rmosolgo/graphql-ruby/compare/35542e3b8c997f52329f90dce9d678add456a3ff..96edf004df4601bad5a610dd463d52595a0f60ed

[rmosolgo]

Thanks for this contribution, it looks great! Out of curiousity, did you update the benchmarks above after removing the pending_keys.dup call? I wonder if that would improve the memory usage of the Dataloader profile...

[glasses618]

Thanks for this contribution, it looks great! Out of curiousity, did you update the benchmarks above after removing the pending_keys.dup call? I wonder if that would improve the memory usage of the Dataloader profile...

Yes, the current benchmark report represents the latest version. The original version of the benchmark, which utilized pending_keys = @pending_keys.dup (35542e3), can be summarized as follows:

In Ruby 2.7.8, it the increases ~0.94% memory and ~2.8% objects. In Ruby 3.0.6, it increases ~0.96% memory and ~2.5% objects.

After observing the memory and objects increases in the original version, I notice that pending_keys = @pending_keys.dup can be removed.

[rmosolgo]

Cool, thanks again for this improvement!