Merge pull request rmosolgo#4994 from igorbelo/fix/4993-fiber-local-variables-are-not-copied-when-using-async-dataloader

Copy Fiber-local variables over to sources when using AsyncDataloader

Author: rmosolgo

lib/graphql/dataloader/async_dataloader.rb

@@ -21,6 +21,7 @@ def run
         manager = spawn_fiber do
           while first_pass || job_fibers.any?
             first_pass = false
+            fiber_vars = get_fiber_variables
 
             while (f = (job_fibers.shift || spawn_job_fiber))
               if f.alive?
@@ -34,6 +35,7 @@ def run
             next_job_fibers.clear
 
             Sync do |root_task|
+              set_fiber_variables(fiber_vars)
               while source_tasks.any? || @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) }
                 while (task = source_tasks.shift || spawn_source_task(root_task, sources_condition))
                   if task.alive?

spec/graphql/dataloader/async_dataloader_spec.rb

@@ -52,6 +52,12 @@ def fetch(keys)
         end
       end
 
+      class FiberLocalContextSource < GraphQL::Dataloader::Source
+        def fetch(keys)
+          keys.map { |key| Thread.current[key] }
+        end
+      end
+
       class Sleeper < GraphQL::Schema::Object
         field :sleeper, Sleeper, null: false, resolver_method: :sleep do
           argument :duration, Float
@@ -126,6 +132,13 @@ def list_waiters(wait:, tags:)
           Kernel.sleep(0.1)
           tags.map { |t| { tag: t, wait: wait }}
         end
+
+        field :fiber_local_context, String do
+          argument :key, String
+        end
+        def fiber_local_context(key:)
+          dataloader.with(FiberLocalContextSource).load(key)
+        end
       end
 
       query(Query)
@@ -273,6 +286,20 @@ def self.included(child_class)
             assert_in_delta 0.3, t2 - t1, 0.06, "Wait was parallel"
             assert_equal [["a", "b", "c"]], AsyncSchema::KeyWaitForSource.fetches, "All keys were fetched at once"
           end
+
+          it 'copies fiber-local variables over to sources' do
+            key = 'arbitrary_context'
+            value = 'test'
+            Thread.current[key] = value
+            query_str = <<-GRAPHQL
+              {
+                fiberLocalContext(key: "#{key}")
+              }
+            GRAPHQL
+
+            result = AsyncSchema.execute(query_str)
+            assert_equal value, result['data']['fiberLocalContext']
+          end
         end
       end
     end