GraphQL::Dataloader::Source#sync complexity

[glasses618]

I noticed that in GraphQL::Dataloader::Source#sync, it is necessary to check all the pending_keys instead of just the key where load was assigned.

To elaborate, the length of pending_keys is 1 when load is initially called in the first sync operation.

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

Lines 37 to 42 in 250a724

	def load(key)
	if @results.key?(key)
	result_for(key)
	else
	@pending_keys << key
	sync

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

Lines 70 to 74 in 250a724

	def sync
	pending_keys = @pending_keys.dup
	@dataloader.yield
	iterations = 0
	while pending_keys.any? { |k| !@results.key?(k) }

In subsequent sync operations, when load is called again, the length of pending_keys increases to 2.
Similarly, for the Nth time load is called, the length of pending_keys becomes N in sync.
As the number of load calls increases, the complexity of the check pending_keys.any? { |k| [email protected]?(k) } also increases.
Is it possible to modify sync so that pending_keys corresponds to the key where load was assigned?

For example:

def load(key)
  if @results.key?(key)
    result_for(key)
  else
    @pending_keys << key
    sync([key])
    result_for(key)
  end
end

def load_all(keys)
 if keys.any? { |k| !@results.key?(k) }
   pending_keys = keys.select { |k| !@results.key?(k) }
   @pending_keys.concat(pending_keys)
   sync(pending_keys)
 end

 keys.map { |k| result_for(k) }
end

def sync(pending_keys)
 pending_keys = pending_keys.dup
 @dataloader.yield
 iterations = 0
 while pending_keys.any? { |k| !@results.key?(k) }
   iterations += 1
   if iterations > 1000
     raise "#{self.class}#sync tried 1000 times to load pending keys (#{pending_keys}), but they still weren't loaded. There is likely a circular dependency."
   end
   @dataloader.yield
 end
 nil
end

[rmosolgo]

Hey, thanks for the great write-up and interesting suggestion. I can see how this would improve the precision of dataloading in theory, but I'm wondering, how does this affect its efficiency in practice?

My guess would have been that a given Source would resolve all of its pending keys in a single fetch(...) call. But did you encounter a scenario where .load waited for more than fetch call to the same Source?

[glasses618]

Thank you for your response.
I have come across a performance issue while using ActiveRecord as the key for GraphQL::Dataloader::Source compared to using GraphQL::Batch::Loader from graphql-batch.

I use StackProf found it takes much time in ActiveRecord::Core#hash which is due to while pending_keys.any? { |k| [email protected]?(k) }.

Here is the result for

query {
  allSameRecords(size: 5000){
    dataloaderUser {
      id
    }
  }
}

(I put the source code at the end)
flamegraph_raw.zip

While I discovered that I can avoid using ActiveRecord as a key, I am currently considering that there might be a bottleneck in the while pending_keys.any? { |k| [email protected]?(k) } that can be improved.

I published a repository to demo this behavior.

graphql version: 2.0.22
rails version: 6.1.7.2
graphql-batch version: 0.5.3

The main code is as follows:

Sources::Association of GraphQL::Dataloader::Source

This is a Source which use ActiceRecord as key

# frozen_string_literal: true

class Sources::Association < GraphQL::Dataloader::Source
  def initialize(model, association_name)
    @model = model
    @association_name = association_name
    validate
  end

  def load(record)
    fail TypeError, "#{@model} loader can't load association for #{record.class}" if !record.is_a?(@model)
    super(record)
  end

  def fetch(records)
    preload_association(records)
    records.map{|s| read_association(s) }
  end

  private

  def validate
    return if @model.reflect_on_association(@association_name)

    fail ArgumentError, "No association #{@association_name} on #{@model}"
  end

  def preload_association(records)
    ::ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
  end

  def read_association(record)
    record.public_send(@association_name)
  end

  def association_loaded?(record)
    record.association(@association_name).loaded?
  end
end

Sources::AssociationV2 of GraphQL::Dataloader::Source

This is a Source which use ActiceRecord's object_id as key

# frozen_string_literal: true

class Sources::AssociationV2 < GraphQL::Dataloader::Source
  def initialize(model, association_name)
    @model = model
    @association_name = association_name
    validate
    @cache = {}
  end

  def load(record)
    fail TypeError, "#{@model} loader can't load association for #{record.class}" if !record.is_a?(@model)
    cache_key = cache_key(record)
    @cache[cache_key] = record
    super(cache_key)
  end

  def fetch(cache_keys)
    records = cache_keys.map{|k| @cache[k] ||= @model.find(k) }
    preload_association(records)
    records.map{|s| read_association(s) }
  end

  private

  def validate
    return if @model.reflect_on_association(@association_name)

    fail ArgumentError, "No association #{@association_name} on #{@model}"
  end

  def preload_association(records)
    ::ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
  end

  def read_association(record)
    record.public_send(@association_name)
  end

  def association_loaded?(record)
    record.association(@association_name).loaded?
  end

  def cache_key(record)
    record.object_id
  end
end

AssociationLoader of GraphQL::Batch::Loader

# frozen_string_literal: true

class AssociationLoader < GraphQL::Batch::Loader
  def self.validate(model, association_name)
    new(model, association_name)
    nil
  end

  def initialize(model, association_name)
    super()
    @model = model
    @association_name = association_name
    validate
  end

  def load(record)
    raise TypeError, "#{@model} loader can't load association for #{record.class}" unless record.is_a?(@model)
    return Promise.resolve(read_association(record)) if association_loaded?(record)
    super
  end

  # We want to load the associations on all records, even if they have the same id
  def cache_key(record)
    record.object_id
  end

  def perform(records)
    preload_association(records)
    records.each { |record| fulfill(record, read_association(record)) }
  end

  private

  def validate
    unless @model.reflect_on_association(@association_name)
      raise ArgumentError, "No association #{@association_name} on #{@model}"
    end
  end

  def preload_association(records)
    ::ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
  end

  def read_association(record)
    record.public_send(@association_name)
  end

  def association_loaded?(record)
    record.association(@association_name).loaded?
  end
end

Query Type

module Types
  class QueryType < Types::BaseObject
    # Add `node(id: ID!) and `nodes(ids: [ID!]!)`
    include GraphQL::Types::Relay::HasNodeField
    include GraphQL::Types::Relay::HasNodesField

    field :all_same_records, [Types::PostType] do
      argument :size, Integer, required: true
    end

    field :all_different_records, [Types::PostType] do
      argument :size, Integer, required: true
    end

    def all_same_records(size:)
      post = Post.joins(:user).last
      Array.new(size){ post }
    end

    def all_different_records(size:)
      Post.joins(:user).last(size)
    end
  end
end

User type

# frozen_string_literal: true

module Types
  class UserType < BaseObject
    field :id, Integer
  end
end

Post Type

# frozen_string_literal: true

module Types
  class PostType < BaseObject
    field :dataloader_user, Types::UserType, null: false
    field :dataloader_user_v2, Types::UserType, null: false
    field :graphql_batch_user, Types::UserType, null: false

    def graphql_batch_user
      AssociationLoader.for(object.class, :user).load(object)
    end

    def dataloader_user
      dataloader.with(::Sources::Association, object.class, :user).load(object)
    end

    def dataloader_user_v2
      dataloader.with(::Sources::AssociationV2, object.class, :user).load(object)
    end
  end
end

Test Result

query {
  allSameRecords(size: 5000){
    dataloaderUser {
      id
    }
  }
}

Before（original ）

After

• use my revised version in GraphQL::Dataloader::Source#sync complexity #4512 (comment)
  which check only the key where load was assigned.

query {
  allSameRecords(size: 5000){
    dataloaderUserV2 {
      id
    }
  }
}

Before

After

query {
  allSameRecords(size: 5000){
    graphqlBatchUser {
      id
    }
  }
}

query {
  allDifferentRecords(size: 5000){
    dataloaderUser {
      id
    }
  }
}

Before

After

query {
  allDifferentRecords(size: 5000){
    dataloaderUserV2 {
      id
    }
  }
}

Before

After

query {
  allDifferentRecords(size: 5000){
    graphqlBatchUser {
      id
    }
  }
}

[rmosolgo]

Oh, I see! So, in this case, it isn't about the dataloading efficiency, but rather about the Ruby/CPU efficiency. Thanks for sharing what you found in this area -- it's really interesting.

I'd like to adopt the changes you proposed here -- could you open a PR with them?

[glasses618]

Thank you. I would open a PR!

[glasses618]

The PR #4519 has been completed. Therefore, I am closing the discussion. Thank you!

[glasses618]

I am uncertain whether I should close the discussion. If I do incorrectly, please remind me. Thank you.