feat: add thread pool with promises to limit concurrent sideloading (#472)

Previously, sideloading in Graphiti used an unbounded number of threads, which could decrease performance due to resource contention and exhaust the ActiveRecord connection pool, leading to timeouts and connection errors.

This commit introduces a thread pool based on ActiveRecord's global_thread_pool_async_query_executor, limiting the maximum number of resources that can be sideloaded concurrently. With a properly configured connection pool, this ensures that the ActiveRecord connection pool is not exhausted by the sideloading process.

Author: MattFenelon

graphiti.gemspec

@@ -22,7 +22,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "jsonapi-renderer", "~> 0.2", ">= 0.2.2"
   spec.add_dependency "dry-types", ">= 0.15.0", "< 2.0"
   spec.add_dependency "graphiti_errors", "~> 1.1.0"
-  spec.add_dependency "concurrent-ruby", "~> 1.0"
+  spec.add_dependency "concurrent-ruby", ">= 1.2", "< 2.0"
   spec.add_dependency "activesupport", ">= 5.2"
 
   spec.add_development_dependency "faraday", "~> 0.15"

lib/graphiti/configuration.rb

@@ -8,6 +8,20 @@ class Configuration
     #   Defaults to false OR if classes are cached (Rails-only)
     attr_accessor :concurrency
 
+    # This number must be considered in accordance with the database
+    # connection pool size configured in `database.yml`. The connection
+    # pool should be large enough to accommodate both the foreground
+    # threads (ie. web server or job worker threads) and background
+    # threads. For each process, Graphiti will create one global
+    # executor that uses this many threads to sideload resources
+    # asynchronously. Thus, the pool size should be at least
+    # `thread_count + concurrency_max_threads + 1`. For example, if your
+    # web server has a maximum of 3 threads, and
+    # `concurrency_max_threads` is set to 4, then your pool size should
+    # be at least 8.
+    # @return [Integer] Maximum number of threads to use when fetching sideloads concurrently
+    attr_accessor :concurrency_max_threads
+
     attr_accessor :respond_to
     attr_accessor :context_for_endpoint
     attr_accessor :links_on_demand
@@ -28,6 +42,7 @@ class Configuration
     def initialize
       @raise_on_missing_sideload = true
       @concurrency = false
+      @concurrency_max_threads = 4
       @respond_to = [:json, :jsonapi, :xml]
       @links_on_demand = false
       @pagination_links_on_demand = false

lib/graphiti/resource_proxy.rb

@@ -77,16 +77,24 @@ def as_graphql(options = {})
     def data
       @data ||= begin
         records = @scope.resolve
-        if records.empty? && raise_on_missing?
-          raise Graphiti::Errors::RecordNotFound
-        end
+        raise Graphiti::Errors::RecordNotFound if records.empty? && raise_on_missing?
+
         records = records[0] if single?
         records
       end
     end
     alias_method :to_a, :data
     alias_method :resolve_data, :data
 
+    def future_resolve_data
+      @scope.future_resolve.then do |records|
+        raise Graphiti::Errors::RecordNotFound if records.empty? && raise_on_missing?
+
+        records = records[0] if single?
+        @data = records
+      end
+    end
+
     def meta
       @meta ||= data.respond_to?(:meta) ? data.meta : {}
     end

lib/graphiti/scope.rb

@@ -2,6 +2,35 @@ module Graphiti
   class Scope
     attr_accessor :object, :unpaginated_object
     attr_reader :pagination
+
+    GLOBAL_THREAD_POOL_EXECUTOR_BROADCAST_STATS = %i[
+      length max_length queue_length max_queue completed_task_count largest_length scheduled_task_count synchronous
+    ]
+    GLOBAL_THREAD_POOL_EXECUTOR = Concurrent::Promises.delay do
+      if Graphiti.config.concurrency
+        concurrency = Graphiti.config.concurrency_max_threads || 4
+        Concurrent::ThreadPoolExecutor.new(
+          min_threads: 0,
+          max_threads: concurrency,
+          max_queue: concurrency * 4,
+          fallback_policy: :caller_runs
+        )
+      else
+        Concurrent::ThreadPoolExecutor.new(max_threads: 0, synchronous: true, fallback_policy: :caller_runs)
+      end
+    end
+    private_constant :GLOBAL_THREAD_POOL_EXECUTOR, :GLOBAL_THREAD_POOL_EXECUTOR_BROADCAST_STATS
+
+    def self.global_thread_pool_executor
+      GLOBAL_THREAD_POOL_EXECUTOR.value!
+    end
+
+    def self.global_thread_pool_stats
+      GLOBAL_THREAD_POOL_EXECUTOR_BROADCAST_STATS.each_with_object({}) do |key, memo|
+        memo[key] = global_thread_pool_executor.send(key)
+      end
+    end
+
     def initialize(object, resource, query, opts = {})
       @object = object
       @resource = resource
@@ -14,57 +43,33 @@ def initialize(object, resource, query, opts = {})
     end
 
     def resolve
-      if @query.zero_results?
-        []
-      else
-        resolved = broadcast_data { |payload|
-          @object = @resource.before_resolve(@object, @query)
-          payload[:results] = @resource.resolve(@object)
-          payload[:results]
-        }
-        resolved.compact!
-        assign_serializer(resolved)
-        yield resolved if block_given?
-        @opts[:after_resolve]&.call(resolved)
-        resolve_sideloads(resolved) unless @query.sideloads.empty?
-        resolved
-      end
+      future_resolve.value!
     end
 
     def resolve_sideloads(results)
-      return if results == []
+      future_resolve_sideloads(results).value!
+    end
 
-      concurrent = Graphiti.config.concurrency
-      promises = []
+    def future_resolve
+      return Concurrent::Promises.fulfilled_future([], self.class.global_thread_pool_executor) if @query.zero_results?
 
-      @query.sideloads.each_pair do |name, q|
-        sideload = @resource.class.sideload(name)
-        next if sideload.nil? || sideload.shared_remote?
-        parent_resource = @resource
-        graphiti_context = Graphiti.context
-        resolve_sideload = -> {
-          Graphiti.config.before_sideload&.call(graphiti_context)
-          Graphiti.context = graphiti_context
-          sideload.resolve(results, q, parent_resource)
-          @resource.adapter.close if concurrent
-        }
-        if concurrent
-          promises << Concurrent::Promise.execute(&resolve_sideload)
-        else
-          resolve_sideload.call
-        end
+      resolved = broadcast_data { |payload|
+        @object = @resource.before_resolve(@object, @query)
+        payload[:results] = @resource.resolve(@object)
+        payload[:results]
+      }
+      resolved.compact!
+      assign_serializer(resolved)
+      yield resolved if block_given?
+      @opts[:after_resolve]&.call(resolved)
+      sideloaded = @query.parents.any?
+      close_adapter = Graphiti.config.concurrency && sideloaded
+      if close_adapter
+        @resource.adapter.close
       end
 
-      if concurrent
-        # Wait for all promises to finish
-        sleep 0.01 until promises.all? { |p| p.fulfilled? || p.rejected? }
-        # Re-raise the error with correct stacktrace
-        # OPTION** to avoid failing here?? if so need serializable patch
-        # to avoid loading data when association not loaded
-        if (rejected = promises.find(&:rejected?))
-          raise rejected.reason
-        end
-      end
+      future_resolve_sideloads(resolved)
+        .then_on(self.class.global_thread_pool_executor, resolved) { resolved }
     end
 
     def parent_resource
@@ -108,6 +113,74 @@ def updated_at
 
     private
 
+    def future_resolve_sideloads(results)
+      return Concurrent::Promises.fulfilled_future(nil, self.class.global_thread_pool_executor) if results == []
+
+      sideload_promises = @query.sideloads.filter_map do |name, q|
+        sideload = @resource.class.sideload(name)
+        next if sideload.nil? || sideload.shared_remote?
+
+        p = future_with_context(results, q, @resource) do |parent_results, sideload_query, parent_resource|
+          Graphiti.config.before_sideload&.call(Graphiti.context)
+          sideload.future_resolve(parent_results, sideload_query, parent_resource)
+        end
+        p.flat
+      end
+
+      Concurrent::Promises.zip_futures_on(self.class.global_thread_pool_executor, *sideload_promises)
+        .rescue_on(self.class.global_thread_pool_executor) do |*reasons|
+          first_error = reasons.find { |r| r.is_a?(Exception) }
+          raise first_error
+        end
+    end
+
+    def future_with_context(*args)
+      thread_storage = Thread.current.keys.each_with_object({}) do |key, memo|
+        memo[key] = Thread.current[key]
+      end
+      fiber_storage =
+        if Fiber.current.respond_to?(:storage)
+          Fiber.current&.storage&.keys&.each_with_object({}) do |key, memo|
+            memo[key] = Fiber[key]
+          end
+        end
+
+      Concurrent::Promises.future_on(
+        self.class.global_thread_pool_executor, Thread.current.object_id, thread_storage, fiber_storage, *args
+      ) do |thread_id, thread_storage, fiber_storage, *args|
+        wrap_in_rails_executor do
+          execution_context_changed = thread_id != Thread.current.object_id
+          if execution_context_changed
+            thread_storage&.keys&.each_with_object(Thread.current) do |key, thread_current|
+              thread_current[key] = thread_storage[key]
+            end
+            fiber_storage&.keys&.each_with_object(Fiber) do |key, fiber_current|
+              fiber_current[key] = fiber_storage[key]
+            end
+          end
+
+          result = Graphiti.broadcast(:global_thread_pool_task_run, self.class.global_thread_pool_stats) do
+            yield(*args)
+          end
+
+          if execution_context_changed
+            thread_storage&.keys&.each { |key| Thread.current[key] = nil }
+            fiber_storage&.keys&.each { |key| Fiber[key] = nil }
+          end
+
+          result
+        end
+      end
+    end
+
+    def wrap_in_rails_executor(&block)
+      if defined?(::Rails.application.executor)
+        ::Rails.application.executor.wrap(&block)
+      else
+        yield
+      end
+    end
+
     def sideload_resource_proxies
       @sideload_resource_proxies ||= begin
         @object = @resource.before_resolve(@object, @query)

lib/graphiti/sideload.rb

@@ -236,7 +236,7 @@ def build_resource_proxy(parents, query, graph_parent)
     end
 
     def load(parents, query, graph_parent)
-      build_resource_proxy(parents, query, graph_parent).to_a
+      future_load(parents, query, graph_parent).value!
     end
 
     # Override in subclass
@@ -286,7 +286,7 @@ def assign(parents, children)
       children.replace(associated) if track_associated
     end
 
-    def resolve(parents, query, graph_parent)
+    def future_resolve(parents, query, graph_parent)
       if single? && parents.length > 1
         raise Errors::SingularSideload.new(self, parents.length)
       end
@@ -300,11 +300,11 @@ def resolve(parents, query, graph_parent)
           sideload: self,
           sideload_parent_length: parents.length,
           default_paginate: false
-        sideload_scope.resolve do |sideload_results|
+        sideload_scope.future_resolve do |sideload_results|
           fire_assign(parents, sideload_results)
         end
       else
-        load(parents, query, graph_parent)
+        future_load(parents, query, graph_parent)
       end
     end
 
@@ -368,6 +368,11 @@ def resource_class_loaded?
 
     private
 
+    def future_load(parents, query, graph_parent)
+      proxy = build_resource_proxy(parents, query, graph_parent)
+      proxy.respond_to?(:future_resolve_data) ? proxy.future_resolve_data : Concurrent::Promises.fulfilled_future(proxy)
+    end
+
     def blank_query?(params)
       if (filter = params[:filter])
         if filter.values == [""]

lib/graphiti/sideload/polymorphic_belongs_to.rb

@@ -108,18 +108,23 @@ def child_for_type!(type)
   end
 
   def resolve(parents, query, graph_parent)
-    parents.group_by(&grouper.field_name).each_pair do |group_name, group|
+    future_resolve(parents, query, graph_parent).value!
+  end
+
+  def future_resolve(parents, query, graph_parent)
+    promises = parents.group_by(&grouper.field_name).filter_map do |(group_name, group)|
       next if group_name.nil? || grouper.ignore?(group_name)
 
       match = ->(c) { c.group_name == group_name.to_sym }
       if (sideload = children.values.find(&match))
         duped = remove_invalid_sideloads(sideload.resource, query)
-        sideload.resolve(group, duped, graph_parent)
+        sideload.future_resolve(group, duped, graph_parent)
       else
         err = ::Graphiti::Errors::PolymorphicSideloadChildNotFound
         raise err.new(self, group_name)
       end
     end
+    Concurrent::Promises.zip(*promises)
   end
 
   private

spec/configuration_spec.rb

@@ -150,6 +150,21 @@
     end
   end
 
+  describe "#concurrency_max_threads" do
+    include_context "with config", :concurrency_max_threads
+
+    it "defaults" do
+      expect(Graphiti.config.concurrency_max_threads).to eq(4)
+    end
+
+    it "is overridable" do
+      Graphiti.configure do |c|
+        c.concurrency_max_threads = 1
+      end
+      expect(Graphiti.config.concurrency_max_threads).to eq(1)
+    end
+  end
+
   describe "#raise_on_missing_sideload" do
     include_context "with config", :raise_on_missing_sideload