fix: patch dataloader to propagate context to new fiber (#1760)

Author: xuanduc987

instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/instrumentation.rb

@@ -27,6 +27,8 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
             require_relative 'tracers/graphql_trace'
             install_new_tracer(config)
           end
+
+          patch
         end
 
         present do
@@ -41,6 +43,10 @@ def supports_new_tracer?
           Gem::Requirement.new('>= 2.0.19').satisfied_by?(gem_version)
         end
 
+        def dataloader_has_spawn_fiber?
+          Gem::Requirement.new('>= 2.1.8').satisfied_by?(gem_version)
+        end
+
         ## Supported configuration keys for the install config hash:
         #
         # The enable_platform_field key expects a boolean value,
@@ -96,6 +102,13 @@ def install_new_tracer(config = {})
             end
           end
         end
+
+        def patch
+          return unless dataloader_has_spawn_fiber?
+
+          require_relative 'patches/dataloader'
+          ::GraphQL::Dataloader.prepend(Patches::Dataloader)
+        end
       end
     end
   end

instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/patches/dataloader.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module OpenTelemetry
+  module Instrumentation
+    module GraphQL
+      module Patches
+        # Patches GraphQL::Dataloader to propagate context to new fiber
+        module Dataloader
+          def spawn_fiber(&block)
+            ctx = OpenTelemetry::Context.current
+            super do
+              OpenTelemetry::Context.with_current(ctx, &block)
+            end
+          end
+        end
+      end
+    end
+  end
+end

instrumentation/graphql/test/instrumentation/graphql/patches/dataloader_test.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require 'test_helper'
+
+require_relative '../../../../lib/opentelemetry/instrumentation/graphql'
+require_relative '../../../../lib/opentelemetry/instrumentation/graphql/patches/dataloader'
+
+describe OpenTelemetry::Instrumentation::GraphQL::Patches::Dataloader do
+  let(:instrumentation) { OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance }
+  let(:tracer) { OpenTelemetry.tracer_provider.tracer('test') }
+  let(:exporter) { EXPORTER }
+  let(:spans) { exporter.finished_spans }
+
+  before do
+    instrumentation.instance_variable_set(:@installed, false)
+    instrumentation.install({})
+  end
+
+  if OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.dataloader_has_spawn_fiber?
+    describe '#spawn_fiber' do
+      it 'set context in the child fiber' do
+        tracer.in_span('parent') do
+          fiber = GraphQL::Dataloader.new.spawn_fiber do
+            tracer.in_span('child1') do
+              # empty block
+            end
+            Fiber.yield
+            tracer.in_span('child2') do
+              # empty block
+            end
+          end
+          fiber.resume
+          fiber.resume
+        end
+
+        parent_span = spans.find { |s| s.name == 'parent' }
+        child1_span = spans.find { |s| s.name == 'child1' }
+        child2_span = spans.find { |s| s.name == 'child2' }
+
+        _(parent_span.span_id).must_equal(child1_span.parent_span_id)
+        _(parent_span.span_id).must_equal(child2_span.parent_span_id)
+      end
+    end
+  end
+end