graphql 2.5.11 → 2.6.10

data/lib/graphql/dataloader/async_dataloader.rb

@@ -1,99 +1,375 @@
 # frozen_string_literal: true
+
 module GraphQL
   class Dataloader
     class AsyncDataloader < Dataloader
+      def self.use(...)
+        install_graphql_methods
+        super
+      end
+
+      def self.install_graphql_methods
+        if !Async::Task.method_defined?(:cancel)
+          Async::Task.alias_method(:cancel, :stop)
+        end
+        if !Async::Task.method_defined?(:graphql_async_dataloader_run)
+          Async::Task.attr_accessor(:graphql_async_dataloader_run)
+          Async::Task.attr_accessor(:graphql_async_dataloader_condition)
+        end
+      end
+
+      def initialize(...)
+        super
+        create_pending_run
+      end
+
+      # @api private
+      attr_reader :pending_sources
+
+      def create_pending_run
+        jobs_fiber_limit, total_fiber_limit = calculate_fiber_limit
+        @pending_run = Run.new(self, total_fiber_limit, jobs_fiber_limit)
+      end
+
       def yield(source = Fiber[:__graphql_current_dataloader_source])
-        trace = Fiber[:__graphql_current_multiplex]&.current_trace
+        task = Async::Task.current
+        run = task.graphql_async_dataloader_run
+        trace = run.trace
         trace&.dataloader_fiber_yield(source)
-        if (condition = Fiber[:graphql_dataloader_next_tick])
-          condition.wait
-        else
-          Fiber.yield
+        if !run.push_task_message(:paused_task, task)
+          task.stop
+        end
+        condition = task.graphql_async_dataloader_condition
+        condition.wait
+        if !run.push_task_message(:resumed_task, task)
+          task.stop
         end
         trace&.dataloader_fiber_resume(source)
         nil
       end
 
-      def run
-        trace = Fiber[:__graphql_current_multiplex]&.current_trace
-        jobs_fiber_limit, total_fiber_limit = calculate_fiber_limit
-        job_fibers = []
-        next_job_fibers = []
-        source_tasks = []
-        next_source_tasks = []
-        first_pass = true
-        sources_condition = Async::Condition.new
-        manager = spawn_fiber do
-          trace&.begin_dataloader(self)
-          while first_pass || !job_fibers.empty?
-            first_pass = false
-            fiber_vars = get_fiber_variables
-
-            while (f = (job_fibers.shift || (((job_fibers.size + next_job_fibers.size + source_tasks.size) < jobs_fiber_limit) && spawn_job_fiber(trace))))
-              if f.alive?
-                finished = run_fiber(f)
-                if !finished
-                  next_job_fibers << f
+      class Run
+        def initialize(dataloader, total_fiber_limit, jobs_fiber_limit)
+          @dataloader = dataloader
+          @root_task = nil
+          @trace = nil
+          @jobs = []
+
+          @total_fiber_limit = total_fiber_limit
+          @jobs_fiber_limit = jobs_fiber_limit
+          @lazies_at_depth = Hash.new { |h, k| h[k] = [] }
+
+          @running_tasks = nil
+          @tasks_channel = nil
+          @tasks_channel_task = nil
+          @activity = nil
+          @task_error = nil
+          @expected_resumes = 0
+          @mode = nil
+
+          @snoozed_jobs_condition = Async::Condition.new
+          @snoozed_sources_condition = Async::Condition.new
+        end
+
+        attr_accessor :trace, :root_task
+
+        attr_reader :dataloader, :jobs, :lazies_at_depth, :jobs_fiber_limit,  :snoozed_jobs_condition, :snoozed_sources_condition
+
+        def jobs_bandwidth?
+          running_count < @jobs_fiber_limit
+        end
+
+        def sources_bandwidth?
+          running_count < current_sources_fiber_limit
+        end
+
+        def close_queues
+          @tasks_channel.close
+          @tasks_channel_task.cancel
+        end
+
+        # Push to the tasks_channel, tolerating a closed channel: on the error path, `run_queue`
+        # closes the channel while sibling tasks can still run one more slice before
+        # `root_task.cancel` reaches them. Record `:task_error` payloads so they aren't lost, and
+        # return false so the caller can stop the task instead of raising `ClosedError` into user code.
+        def push_task_message(msg, data)
+          @tasks_channel.push([msg, data])
+          true
+        rescue Async::Queue::ClosedError
+          if msg == :task_error
+            @task_error ||= data
+          end
+          false
+        end
+
+        def wait_for_activity
+          @activity.wait
+        end
+
+        def quiesced?
+          @running_tasks.empty? && @tasks_channel.empty? && @expected_resumes == 0
+        end
+
+        def has_pending_work?
+          @mode == :jobs ? @jobs.any? : @dataloader.pending_sources.any?(&:pending?) # rubocop:disable Development/NoneWithoutBlockCop
+        end
+
+        def has_bandwidth?
+          @mode == :jobs ? jobs_bandwidth? : sources_bandwidth?
+        end
+
+        # Signalled tasks don't appear in any accounting until their first slice
+        # pushes `:resumed_task`, so they have to be counted at signal time:
+        def expect_resumes(count)
+          @expected_resumes = count
+        end
+
+        def check_error!
+          if (err = @task_error)
+            @task_error = nil
+            raise err
+          end
+        end
+
+        def new_queues(mode)
+          @mode = mode
+          @tasks_channel = Async::Queue.new(parent: @root_task)
+          @activity = Async::Condition.new
+          @task_error = nil
+          @expected_resumes = 0
+          @running_tasks = []
+          @tasks_channel_task = @root_task.async do |_t|
+            while ((msg, data) = @tasks_channel.wait)
+              case msg
+              when :started_task
+                @running_tasks.push(data)
+                data.run
+              when :resumed_task
+                if @expected_resumes > 0
+                  @expected_resumes -= 1
                 end
+                @running_tasks.push(data)
+              when :finished_task, :paused_task
+                @running_tasks.delete(data)
+              when :task_error
+                @task_error ||= data
+              else
+                raise ArgumentError, "Unknown tasks_channel action: #{msg.inspect}"
               end
+              @activity.signal
+            end
+          end
+        end
+
+        def running?
+          @snoozed_jobs_condition.waiting? || @snoozed_sources_condition.waiting?
+        end
+
+        def current_sources_fiber_limit
+          within_limit = @total_fiber_limit - running_count
+          if within_limit < 1
+            1
+          else
+            within_limit
+          end
+        end
+
+        private
+
+        def running_count
+          @snoozed_jobs_condition.instance_variable_get(:@ready).num_waiting +
+            @snoozed_sources_condition.instance_variable_get(:@ready).num_waiting +
+            (@running_tasks&.size || 0)
+        end
+      end
+
+      def append_job(callable = nil, &block)
+        active_run.jobs.push(callable || block)
+        nil
+      end
+
+      def lazy_at_depth(depth, lazy)
+        active_run.lazies_at_depth[depth] << lazy
+      end
+
+      def active_run
+        @pending_run || current_task_run || raise(GraphQL::Error, "No available Run to append to, GraphQL-Ruby bug")
+      end
+
+      # The current task's run, but only if it belongs to this dataloader. A different
+      # dataloader may be running inside one of our tasks (or vice versa), e.g. a query
+      # executed from a resolver or a subscription trigger; its run must not be reused.
+      def current_task_run
+        run = Async::Task.current?&.graphql_async_dataloader_run
+        run if run&.dataloader.equal?(self)
+      end
+
+      def run_isolated
+        previous_run = current_task_run
+        prev_pending_keys = {}
+        # Clear pending loads but keep already-cached records
+        # in case they are useful to the given block.
+        @source_cache.each do |source_class, batched_sources|
+          batched_sources.each do |batch_args, batched_source_instance|
+            if batched_source_instance.pending?
+              prev_pending_keys[batched_source_instance] = batched_source_instance.pending.dup
+              batched_source_instance.pending.clear
             end
-            job_fibers.concat(next_job_fibers)
-            next_job_fibers.clear
-
-            Sync do |root_task|
-              set_fiber_variables(fiber_vars)
-              while !source_tasks.empty? || @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) }
-                while (task = (source_tasks.shift || (((job_fibers.size + next_job_fibers.size + source_tasks.size + next_source_tasks.size) < total_fiber_limit) && spawn_source_task(root_task, sources_condition, trace))))
-                  if task.alive?
-                    root_task.yield # give the source task a chance to run
-                    next_source_tasks << task
+          end
+        end
+
+        res = nil
+        create_pending_run
+        @pending_run.jobs << -> { res = yield }
+        run
+        res
+      ensure
+        if previous_run
+          Async::Task.current.graphql_async_dataloader_run = previous_run
+          # clear the one created in #run:
+          @pending_run = nil
+        end
+        prev_pending_keys.each do |source_instance, pending|
+          pending.each do |key, value|
+            next if source_instance.results.key?(key)
+
+            queue_pending_source(source_instance) if source_instance.pending.empty?
+            source_instance.pending[key] = value
+          end
+        end
+      end
+
+      def run(trace_query_lazy: nil)
+        trace = Fiber[:__graphql_current_multiplex]&.current_trace
+        run = @pending_run || current_task_run || raise(GraphQL::Error, "No available Run, GraphQL-Ruby internal bug")
+        @pending_run = nil
+        run.trace = trace
+        first_pass = true
+        trace&.begin_dataloader(self)
+        fiber_vars = get_fiber_variables
+        raised_error = nil
+        jobs = run.jobs
+        Sync do |_maybe_new_task|
+          # Make sure there's a new task instance to hold `.graphql_...` state:
+          task = Async::Task.new do |root_task|
+            run.root_task = root_task
+            root_task.graphql_async_dataloader_run = run
+            set_fiber_variables(fiber_vars)
+
+            while first_pass || run.running? || !jobs.empty?
+              first_pass = false
+              run_queue(run, run.snoozed_jobs_condition, :jobs)
+              run_queue(run, run.snoozed_sources_condition, :sources)
+
+              if !run.lazies_at_depth.empty?
+                with_trace_query_lazy(trace_query_lazy) do
+                  if enqueue_next_pending_lazies(run.lazies_at_depth)
+                    run_queue(run, run.snoozed_jobs_condition, :jobs)
                   end
                 end
-                sources_condition.signal
-                source_tasks.concat(next_source_tasks)
-                next_source_tasks.clear
               end
             end
+          rescue StandardError => err
+            raised_error = err
+            root_task.cancel
           end
-          trace&.end_dataloader(self)
-        end
 
-        manager.resume
-        if manager.alive?
-          raise "Invariant: Manager didn't terminate successfully: #{manager}"
+          task.run
+          task.wait
         end
-
+        create_pending_run
+        if raised_error
+          raise raised_error
+        end
+        trace&.end_dataloader(self)
       rescue UncaughtThrowError => e
         throw e.tag, e.value
       end
 
       private
 
-      def spawn_source_task(parent_task, condition, trace)
-        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
+      def run_queue(run, condition, mode)
+        opened_queues = false
+
+        if condition.waiting?
+          opened_queues = true
+          run.new_queues(mode)
+          run.expect_resumes(condition.instance_variable_get(:@ready).num_waiting)
+          condition.signal
+        end
+
+        loop do
+          pending_work = (mode == :jobs) ? (!run.jobs.empty? && run.jobs_bandwidth? ? run.jobs : nil) : (drain_pending_sources)
+          if pending_work
+            if opened_queues == false
+              opened_queues = true
+              run.new_queues(mode)
+            end
+            num_tasks = mode == :sources ? run.current_sources_fiber_limit : 1
+            if num_tasks > pending_work.size
+              num_tasks = pending_work.size
+            end
+            spawn_tasks(run, mode, condition, pending_work, num_tasks)
+          end
+
+          if !opened_queues
+            break
+          end
+
+          run.check_error!
+
+          if run.quiesced?
+            if !run.has_pending_work? || !run.has_bandwidth?
+              break
             end
+            # Quiesced, but more work appeared - loop around to drain it.
+          else
+            run.wait_for_activity
           end
         end
+      ensure
+        if opened_queues
+          run.close_queues
+        end
+      end
 
-        if pending_sources
-          fiber_vars = get_fiber_variables
-          parent_task.async do
-            trace&.dataloader_spawn_source_fiber(pending_sources)
+      # Use a separate method for this so that the outer loop's reassignment of `pending_work`
+      # doesn't affect already-running tasks which (would) close over that variable
+      def spawn_tasks(run, mode, condition, pending_work, num_tasks)
+        fiber_vars = get_fiber_variables
+        trace = run.trace
+        num_tasks.times do
+          new_task = Async::Task.new(run.root_task) do |task|
+            task.graphql_async_dataloader_run = run
+            task.graphql_async_dataloader_condition = condition
             set_fiber_variables(fiber_vars)
-            Fiber[:graphql_dataloader_next_tick] = condition
-            pending_sources.each do |s|
-              trace&.begin_dataloader_source(s)
-              s.run_pending_keys
-              trace&.end_dataloader_source(s)
+            case mode
+            when :jobs
+              trace&.dataloader_spawn_execution_fiber(pending_work)
+              while job = pending_work.shift
+                job.call
+              end
+            when :sources
+              trace&.dataloader_spawn_source_fiber(pending_work)
+              while (source = pending_work.shift)
+                Fiber[:__graphql_current_dataloader_source] = source
+                trace&.begin_dataloader_source(source)
+                source.run_pending_keys
+                trace&.end_dataloader_source(source)
+              end
+            else
+              raise ArgumentError, "Unknown mode: #{mode.inspect}"
             end
+            nil
+          rescue StandardError => err
+            run.push_task_message(:task_error, err)
+          else
+            run.push_task_message(:finished_task, task)
+          ensure
             cleanup_fiber
             trace&.dataloader_fiber_exit
           end
+          run.push_task_message(:started_task, new_task)
         end
       end
     end

data/lib/graphql/dataloader/null_dataloader.rb

@@ -2,24 +2,62 @@
 
 module GraphQL
   class Dataloader
-    # The default implementation of dataloading -- all no-ops.
+    # GraphQL-Ruby uses this when Dataloader isn't enabled.
     #
-    # The Dataloader interface isn't public, but it enables
-    # simple internal code while adding the option to add Dataloader.
+    # It runs execution code inline and gathers lazy objects (eg. Promises)
+    # and resolves them during {#run}.
     class NullDataloader < Dataloader
-      # These are all no-ops because code was
-      # executed synchronously.
+      def initialize(*)
+        @lazies_at_depth = Hash.new { |h,k| h[k] = [] }
+      end
+
+      def freeze
+        @lazies_at_depth.default_proc = nil
+        @lazies_at_depth.freeze
+        super
+      end
+
+      def run(trace_query_lazy: nil)
+        with_trace_query_lazy(trace_query_lazy) do
+          while !@lazies_at_depth.empty?
+            smallest_depth = nil
+            @lazies_at_depth.each_key do |depth_key|
+              smallest_depth ||= depth_key
+              if depth_key < smallest_depth
+                smallest_depth = depth_key
+              end
+            end
+
+            if smallest_depth
+              lazies = @lazies_at_depth.delete(smallest_depth)
+              lazies.each(&:value) # resolve these Lazy instances
+            end
+          end
+        end
+      end
+
+      def run_isolated
+        # Reuse this instance because execution code may already have a reference to _this_ `dataloader` inside the given block.
+        prev_lazies_at_depth = @lazies_at_depth
+        @lazies_at_depth = @lazies_at_depth.dup.clear
+        res = nil
+        append_job {
+          res = yield
+        }
+        run
+        res
+      ensure
+        @lazies_at_depth = prev_lazies_at_depth
+      end
 
-      def initialize(*); end
-      def run; end
-      def run_isolated; yield; end
       def clear_cache; end
+
       def yield(_source)
         raise GraphQL::Error, "GraphQL::Dataloader is not running -- add `use GraphQL::Dataloader` to your schema to use Dataloader sources."
       end
 
-      def append_job
-        yield
+      def append_job(callable = nil)
+        callable ? callable.call : yield
         nil
       end
 

data/lib/graphql/dataloader/source.rb

@@ -20,9 +20,7 @@ module GraphQL
       # @return [Dataloader::Request] a pending request for a value from `key`. Call `.load` on that object to wait for the result.
       def request(value)
         res_key = result_key_for(value)
-        if !@results.key?(res_key)
-          @pending[res_key] ||= normalize_fetch_key(value)
-        end
+        add_pending_key(res_key, value)
         Dataloader::Request.new(self, value)
       end
 
@@ -51,9 +49,7 @@ module GraphQL
       def request_all(values)
         values.each do |v|
           res_key = result_key_for(v)
-          if !@results.key?(res_key)
-            @pending[res_key] ||= normalize_fetch_key(v)
-          end
+          add_pending_key(res_key, v)
         end
         Dataloader::RequestAll.new(self, values)
       end
@@ -65,7 +61,7 @@ module GraphQL
         if @results.key?(result_key)
           result_for(result_key)
         else
-          @pending[result_key] ||= normalize_fetch_key(value)
+          add_pending_key(result_key, value)
           sync([result_key])
           result_for(result_key)
         end
@@ -79,8 +75,7 @@ module GraphQL
         values.each { |v|
           k = result_key_for(v)
           result_keys << k
-          if !@results.key?(k)
-            @pending[k] ||= normalize_fetch_key(v)
+          if add_pending_key(k, v)
             pending_keys << k
           end
         }
@@ -89,7 +84,8 @@ module GraphQL
           sync(pending_keys)
         end
 
-        result_keys.map { |k| result_for(k) }
+        result_keys.map! { |k| result_for(k) }
+        result_keys
       end
 
       # Subclasses must implement this method to return a value for each of `keys`
@@ -105,6 +101,7 @@ module GraphQL
       # Then run the batch and update the cache.
       # @return [void]
       def sync(pending_result_keys)
+        @dataloader.queue_pending_source(self) if pending?
         @dataloader.yield(self)
         iterations = 0
         while pending_result_keys.any? { |key| !@results.key?(key) }
@@ -138,22 +135,25 @@ module GraphQL
       # @api private
       # @return [void]
       def run_pending_keys
-        if !@fetching.empty?
-          @fetching.each_key { |k| @pending.delete(k) }
-        end
+        @fetching.each_key { |k| @pending.delete(k) }
         return if @pending.empty?
         fetch_h = @pending
-        @pending = {}
         @fetching.merge!(fetch_h)
+        @pending = {}
         results = fetch(fetch_h.values)
-        fetch_h.each_with_index do |(key, _value), idx|
+        idx = 0
+
+        fetch_h.each_key do |key|
           @results[key] = results[idx]
+          @fetching.delete(key)
+          idx += 1
         end
         nil
       rescue StandardError => error
-        fetch_h.each_key { |key| @results[key] = error }
-      ensure
-        fetch_h && fetch_h.each_key { |k| @fetching.delete(k) }
+        fetch_h.each_key { |key|
+          @results[key] = error
+          @fetching.delete(key)
+        }
       end
 
       # These arguments are given to `dataloader.with(source_class, ...)`. The object
@@ -171,7 +171,11 @@ module GraphQL
       # @param batch_kwargs [Hash]
       # @return [Object]
       def self.batch_key_for(*batch_args, **batch_kwargs)
-        [*batch_args, **batch_kwargs]
+        if batch_kwargs.any? # rubocop:disable Development/NoneWithoutBlockCop
+          [*batch_args, **batch_kwargs]
+        else
+          batch_args
+        end
       end
 
       # Clear any already-loaded objects for this source
@@ -185,6 +189,15 @@ module GraphQL
 
       private
 
+      def add_pending_key(result_key, value)
+        return false if @results.key?(result_key)
+
+        was_empty = @pending.empty?
+        @pending[result_key] ||= normalize_fetch_key(value)
+        @dataloader.queue_pending_source(self) if was_empty
+        true
+      end
+
       # Reads and returns the result for the key from the internal cache, or raises an error if the result was an error
       # @param key [Object] key passed to {#load} or {#load_all}
       # @return [Object] The result from {#fetch} for `key`.