Rewrite NotificationsTrace

.rubocop.yml

@@ -48,7 +48,7 @@ Development/NoFocusCop:
 Development/TraceMethodsCop:
   Include:
     - "lib/graphql/tracing/perfetto_trace.rb"
-    - "lib/graphql/tracing/new_relic_trace.rb"
+    - "lib/graphql/tracing/notifications_trace.rb"
 
 # def ...
 # end

cop/development/trace_methods_cop.rb

@@ -16,8 +16,6 @@ class TraceMethodsCop < RuboCop::Cop::Base
         :begin_dataloader,
         :begin_dataloader_source,
         :begin_execute_field,
-        :begin_execute_multiplex,
-        :begin_parse,
         :begin_resolve_type,
         :begin_validate,
         :dataloader_fiber_exit,
@@ -30,8 +28,6 @@ class TraceMethodsCop < RuboCop::Cop::Base
         :end_dataloader,
         :end_dataloader_source,
         :end_execute_field,
-        :end_execute_multiplex,
-        :end_parse,
         :end_resolve_type,
         :end_validate,
         :execute_field,
@@ -73,13 +69,20 @@ def on_module(node)
             # Not really necessary for making a good trace:
             :lex, :analyze_query, :execute_query, :execute_query_lazy,
             # Only useful for isolated event tracking:
+            :begin_dataloader, :end_dataloader,
             :dataloader_fiber_exit, :dataloader_spawn_execution_fiber, :dataloader_spawn_source_fiber
           ]
           missing_defs.each do |missing_def|
             if all_defs.include?(:"begin_#{missing_def}") && all_defs.include?(:"end_#{missing_def}")
               redundant_defs << missing_def
               redundant_defs << :"#{missing_def}_lazy"
             end
+            missing_name = missing_def.to_s
+            if missing_name.start_with?("begin") && all_defs.include?(:"#{missing_name.sub("begin_", "")}")
+              redundant_defs << missing_def
+            elsif missing_name.start_with?("end") && all_defs.include?(:"#{missing_name.sub("end_", "")}")
+              redundant_defs << missing_def
+            end
           end
 
           missing_defs -= redundant_defs

lib/graphql/execution/interpreter.rb

@@ -37,7 +37,6 @@ def run_all(schema, query_options, context: {}, max_complexity: schema.max_compl
           multiplex = Execution::Multiplex.new(schema: schema, queries: queries, context: context, max_complexity: max_complexity)
           Fiber[:__graphql_current_multiplex] = multiplex
           trace = multiplex.current_trace
-          trace.begin_execute_multiplex(multiplex)
           trace.execute_multiplex(multiplex: multiplex) do
             schema = multiplex.schema
             queries = multiplex.queries
@@ -154,8 +153,6 @@ def run_all(schema, query_options, context: {}, max_complexity: schema.max_compl
               }
             end
           end
-        ensure
-          trace&.end_execute_multiplex(multiplex)
         end
       end
 

lib/graphql/query.rb

@@ -441,16 +441,13 @@ def prepare_ast
       @warden ||= @schema.warden_class.new(schema: @schema, context: @context)
       parse_error = nil
       @document ||= begin
-        current_trace.begin_parse(query_string)
         if query_string
           GraphQL.parse(query_string, trace: self.current_trace, max_tokens: @schema.max_query_string_tokens)
         end
       rescue GraphQL::ParseError => err
         parse_error = err
         @schema.parse_error(err, @context)
         nil
-      ensure
-        current_trace.end_parse(query_string)
       end
 
       @fragments = {}

lib/graphql/tracing.rb

@@ -25,6 +25,7 @@ module Tracing
     autoload :AppOpticsTrace, "graphql/tracing/appoptics_trace"
     autoload :AppsignalTrace, "graphql/tracing/appsignal_trace"
     autoload :DataDogTrace, "graphql/tracing/data_dog_trace"
+    autoload :MonitorTrace, "graphql/tracing/monitor_trace"
     autoload :NewRelicTrace, "graphql/tracing/new_relic_trace"
     autoload :NotificationsTrace, "graphql/tracing/notifications_trace"
     autoload :SentryTrace, "graphql/tracing/sentry_trace"

lib/graphql/tracing/active_support_notifications_trace.rb

@@ -10,6 +10,13 @@ module Tracing
     #   class MySchema < GraphQL::Schema
     #     trace_with(GraphQL::Tracing::ActiveSupportNotificationsTrace)
     #   end
+    #
+    # @example Subscribing to GraphQL events with ActiveSupport::Notifications
+    #   ActiveSupport::Notifications.subscribe(/graphql/) do |event|
+    #     pp event.name
+    #     pp event.payload
+    #   end
+    #
     module ActiveSupportNotificationsTrace
       include NotificationsTrace
       def initialize(engine: ActiveSupport::Notifications, **rest)

lib/graphql/tracing/monitor_trace.rb

@@ -0,0 +1,228 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Tracing
+    # This module is the basis for Ruby-level integration with third-party monitoring platforms.
+    # Platform-specific traces include this module and implement an adapter.
+    #
+    # @see ActiveSupportNotificationsTrace Integration via ActiveSupport::Notifications, an alternative approach.
+    module MonitorTrace
+      class Monitor
+        def initialize(set_transaction_name:)
+          @set_transaction_name = set_transaction_name
+          @platform_field_key_cache = Hash.new { |h, k| h[k] = platform_field_key(k) }.compare_by_identity
+          @platform_authorized_key_cache = Hash.new { |h, k| h[k] = platform_authorized_key(k) }.compare_by_identity
+          @platform_resolve_type_key_cache = Hash.new { |h, k| h[k] = platform_resolve_type_key(k) }.compare_by_identity
+          @platform_source_class_key_cache = Hash.new { |h, source_cls| h[source_cls] = platform_source_class_key(source_cls) }.compare_by_identity
+        end
+
+        def instrument(keyword, object, &block)
+          raise "Implement #{self.class}#instrument to measure the block"
+        end
+
+        def start_event(keyword, object)
+          ev = self.class::Event.new(self, keyword, object)
+          ev.start
+          ev
+        end
+
+        # Get the transaction name based on the operation type and name if possible, or fall back to a user provided
+        # one. Useful for anonymous queries.
+        def transaction_name(query)
+          selected_op = query.selected_operation
+          txn_name = if selected_op
+            op_type = selected_op.operation_type
+            op_name = selected_op.name || fallback_transaction_name(query.context) || "anonymous"
+            "#{op_type}.#{op_name}"
+          else
+            "query.anonymous"
+          end
+          "GraphQL/#{txn_name}"
+        end
+
+        def fallback_transaction_name(context)
+          context[:tracing_fallback_transaction_name]
+        end
+
+        def name_for(keyword, object)
+          case keyword
+          when :execute_field
+            @platform_field_key_cache[object]
+          when :authorized
+            @platform_authorized_key_cache[object]
+          when :resolve_type
+            @platform_resolve_type_key_cache[object]
+          when :dataloader_source
+            @platform_source_class_key_cache[object.class]
+          when :parse then self.class::PARSE_NAME
+          when :lex then self.class::LEX_NAME
+          when :execute then self.class::EXECUTE_NAME
+          when :analyze then self.class::ANALYZE_NAME
+          when :validate then self.class::VALIDATE_NAME
+          else
+            raise "No name for #{keyword.inspect}"
+          end
+        end
+
+        class Event
+          def initialize(engine, keyword, object)
+            @engine = engine
+            @keyword = keyword
+            @object = object
+          end
+
+          attr_reader :keyword, :object
+
+          def start
+            raise "Implement #{self.class}#start to begin a new event (#{inspect})"
+          end
+
+          def finish
+            raise "Implement #{self.class}#finish to end this event (#{inspect})"
+          end
+        end
+      end
+
+      def self.create_module(monitor_name)
+        if !monitor_name.match?(/[a-z]+/)
+          raise ArgumentError, "monitor name must be [a-z]+, not: #{monitor_name.inspect}"
+        end
+
+        trace_module = Module.new
+        code = MODULE_TEMPLATE % {
+          monitor: monitor_name,
+          monitor_class: monitor_name.capitalize + "Monitor",
+        }
+        trace_module.module_eval(code, __FILE__, __LINE__ + 5)
+        trace_module
+      end
+
+      MODULE_TEMPLATE = <<~RUBY
+        # @param set_transaction_name [Boolean] If `true`, use the GraphQL operation name as the request name on the monitoring platform
+        # @param trace_scalars [Boolean] If `true`, leaf fields will be traced too (Scalars _and_ Enums)
+        # @param trace_authorized [Boolean] If `false`, skip tracing `authorized?` calls
+        # @param trace_resolve_type [Boolean] If `false`, skip tracing `resolve_type?` calls
+        def initialize(set_transaction_name: false, trace_scalars: false, trace_authorized: true, trace_resolve_type: true, **rest)
+          @trace_scalars = trace_scalars
+          @trace_authorized = trace_authorized
+          @trace_resolve_type = trace_resolve_type
+          @set_transaction_name = set_transaction_name
+          @%{monitor} = %{monitor_class}.new(set_transaction_name: @set_transaction_name)
+          super
+        end
+
+        def parse(query_string:)
+          @%{monitor}.instrument(:parse, query_string) do
+            super
+          end
+        end
+
+        def lex(query_string:)
+          @%{monitor}.instrument(:lex, query_string) do
+            super
+          end
+        end
+
+        def validate(query:, validate:)
+          @%{monitor}.instrument(:validate, query) do
+            super
+          end
+        end
+
+        def begin_analyze_multiplex(multiplex, analyzers)
+          begin_%{monitor}_event(:analyze, nil)
+          super
+        end
+
+        def end_analyze_multiplex(multiplex, analyzers)
+          finish_%{monitor}_event
+          super
+        end
+
+        def execute_multiplex(multiplex:)
+          @%{monitor}.instrument(:execute, multiplex) do
+            super
+          end
+        end
+
+        def begin_execute_field(field, object, arguments, query)
+          return_type = field.type.unwrap
+          trace_field = if return_type.kind.scalar? || return_type.kind.enum?
+            (field.trace.nil? && @trace_scalars) || field.trace
+          else
+            true
+          end
+
+          if trace_field
+            begin_%{monitor}_event(:execute_field, field)
+          end
+          super
+        end
+
+        def end_execute_field(field, object, arguments, query, result)
+          finish_%{monitor}_event
+          super
+        end
+
+        def dataloader_fiber_yield(source)
+          Fiber[PREVIOUS_EV_KEY] = finish_%{monitor}_event
+          super
+        end
+
+        def dataloader_fiber_resume(source)
+          prev_ev = Fiber[PREVIOUS_EV_KEY]
+          begin_%{monitor}_event(prev_ev.keyword, prev_ev.object)
+          super
+        end
+
+        def begin_authorized(type, object, context)
+          @trace_authorized && begin_%{monitor}_event(:authorized, type)
+          super
+        end
+
+        def end_authorized(type, object, context, result)
+          finish_%{monitor}_event
+          super
+        end
+
+        def begin_resolve_type(type, value, context)
+          @trace_resolve_type && begin_%{monitor}_event(:resolve_type, type)
+          super
+        end
+
+        def end_resolve_type(type, value, context, resolved_type)
+          finish_%{monitor}_event
+          super
+        end
+
+        def begin_dataloader_source(source)
+          begin_%{monitor}_event(:dataloader_source, source)
+          super
+        end
+
+        def end_dataloader_source(source)
+          finish_%{monitor}_event
+          super
+        end
+
+        CURRENT_EV_KEY = :__graphql_%{monitor}_trace_event
+        PREVIOUS_EV_KEY = :__graphql_%{monitor}_trace_previous_event
+
+        private
+
+        def begin_%{monitor}_event(keyword, object)
+          Fiber[CURRENT_EV_KEY] = @%{monitor}.start_event(keyword, object)
+        end
+
+        def finish_%{monitor}_event
+          if ev = Fiber[CURRENT_EV_KEY]
+            ev.finish
+            # Use `false` to prevent grabbing an event from a parent fiber
+            Fiber[CURRENT_EV_KEY] = false
+            ev
+          end
+        end
+      RUBY
+    end
+  end
+end