Rewrite NewRelicTrace to support fiber stops/starts

.rubocop.yml

@@ -4,7 +4,7 @@ require:
   - ./lib/graphql/rubocop/graphql/default_null_true
   - ./lib/graphql/rubocop/graphql/default_required_true
   - ./cop/development/context_is_passed_cop
-  - ./cop/development/trace_calls_super_cop.rb
+  - ./cop/development/trace_methods_cop.rb
 
 AllCops:
   DisabledByDefault: true
@@ -45,9 +45,10 @@ Development/NoFocusCop:
   Include:
     - "spec/**/*"
 
-Development/TraceCallsSuperCop:
+Development/TraceMethodsCop:
   Include:
-    - "lib/graphql/tracing/*_trace.rb"
+    - "lib/graphql/tracing/perfetto_trace.rb"
+    - "lib/graphql/tracing/new_relic_trace.rb"
 
 # def ...
 # end

cop/development/trace_calls_super_cop.rb → cop/development/trace_methods_cop.rb

@@ -3,7 +3,7 @@
 
 module Cop
   module Development
-    class TraceCallsSuperCop < RuboCop::Cop::Base
+    class TraceMethodsCop < RuboCop::Cop::Base
       extend RuboCop::Cop::AutoCorrector
 
       TRACE_HOOKS = [
@@ -16,7 +16,7 @@ class TraceCallsSuperCop < RuboCop::Cop::Base
         :begin_dataloader,
         :begin_dataloader_source,
         :begin_execute_field,
-        :begin_multiplex,
+        :begin_execute_multiplex,
         :begin_parse,
         :begin_resolve_type,
         :begin_validate,
@@ -30,7 +30,7 @@ class TraceCallsSuperCop < RuboCop::Cop::Base
         :end_dataloader,
         :end_dataloader_source,
         :end_execute_field,
-        :end_multiplex,
+        :end_execute_multiplex,
         :end_parse,
         :end_resolve_type,
         :end_validate,
@@ -51,8 +51,35 @@ class TraceCallsSuperCop < RuboCop::Cop::Base
       def on_def(node)
         if TRACE_HOOKS.include?(node.method_name) && !node.each_descendant(:super, :zsuper).any?
           add_offense(node) do |corrector|
-            offset = node.loc.column + 2
-            corrector.insert_after(node.body.loc.expression, "\n#{' ' * offset}super")
+            if node.body
+              offset = node.loc.column + 2
+              corrector.insert_after(node.body.loc.expression, "\n#{' ' * offset}super")
+            end
+          end
+        end
+      end
+
+      def on_module(node)
+        if node.defined_module_name.to_s.end_with?("Trace")
+          all_defs = []
+          node.body.each_child_node do |body_node|
+            if body_node.def_type?
+              all_defs << body_node.method_name
+            end
+          end
+
+          missing_defs = TRACE_HOOKS - all_defs
+          redundant_defs = [:lex, :analyze_query, :execute_query, :execute_query_lazy]
+          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
+          end
+
+          missing_defs -= redundant_defs
+          if missing_defs.any?
+            add_offense(node, message: "Missing some trace hook methods:\n\n- #{missing_defs.join("\n- ")}")
           end
         end
       end

lib/graphql/execution/interpreter.rb

@@ -37,7 +37,7 @@ 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_multiplex(multiplex)
+          trace.begin_execute_multiplex(multiplex)
           trace.execute_multiplex(multiplex: multiplex) do
             schema = multiplex.schema
             queries = multiplex.queries
@@ -155,7 +155,7 @@ def run_all(schema, query_options, context: {}, max_complexity: schema.max_compl
             end
           end
         ensure
-          trace&.end_multiplex(multiplex)
+          trace&.end_execute_multiplex(multiplex)
         end
       end
 

lib/graphql/static_validation/validator.rb

@@ -27,8 +27,8 @@ def initialize(schema:, rules: GraphQL::StaticValidation::ALL_RULES)
       # @param max_errors [Integer] Maximum number of errors before aborting validation. Any positive number will limit the number of errors. Defaults to nil for no limit.
       # @return [Array<Hash>]
       def validate(query, validate: true, timeout: nil, max_errors: nil)
+        errors = nil
         query.current_trace.begin_validate(query, validate)
-        is_valid = false
         query.current_trace.validate(validate: validate, query: query) do
           begin_t = Time.now
           errors = if validate == false
@@ -53,21 +53,20 @@ def validate(query, validate: true, timeout: nil, max_errors: nil)
 
             context.errors
           end
-          is_valid = errors.size == 0
 
           {
             remaining_timeout: timeout ? (timeout - (Time.now - begin_t)) : nil,
             errors: errors,
           }
         end
       rescue GraphQL::ExecutionError => e
-        is_valid = false
+        errors = [e]
         {
           remaining_timeout: nil,
-          errors: [e],
+          errors: errors,
         }
       ensure
-        query.current_trace.end_validate(query, validate, is_valid)
+        query.current_trace.end_validate(query, validate, errors)
       end
 
       # Invoked when static validation times out.

lib/graphql/tracing/new_relic_trace.rb

@@ -5,72 +5,159 @@
 module GraphQL
   module Tracing
     module NewRelicTrace
-      include PlatformTrace
-
       # @param set_transaction_name [Boolean] If true, the GraphQL operation name will be used as the transaction name.
       #   This is not advised if you run more than one query per HTTP request, for example, with `graphql-client` or multiplexing.
       #   It can also be specified per-query with `context[:set_new_relic_transaction_name]`.
       def initialize(set_transaction_name: false, **_rest)
         @set_transaction_name = set_transaction_name
+        @nr_field_names = Hash.new do |h, field|
+          h[field] = "GraphQL/#{field.owner.graphql_name}/#{field.graphql_name}"
+        end.compare_by_identity
+
+        @nr_authorized_names = Hash.new do |h, type|
+          h[type] = "GraphQL/Authorized/#{type.graphql_name}"
+        end.compare_by_identity
+
+        @nr_resolve_type_names = Hash.new do |h, type|
+          h[type] = "GraphQL/ResolveType/#{type.graphql_name}"
+        end.compare_by_identity
+
+        @nr_source_names = Hash.new do |h, source_inst|
+          h[source_inst] = "GraphQL/Source/#{source_inst.class.name}"
+        end.compare_by_identity
+
+        @nr_parse = @nr_validate = @nr_analyze = @nr_execute = nil
+        super
+      end
+
+      def begin_parse(query_str)
+        @nr_parse = NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: "GraphQL/parse", category: :web)
+        super
+      end
+
+      def end_parse(query_str)
+        @nr_parse.finish
+        super
+      end
+
+      def begin_validate(query, validate)
+        @nr_validate = NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: "GraphQL/validate", category: :web)
+        super
+      end
+
+      def end_validate(query, validate, validation_errors)
+        @nr_validate.finish
+        super
+      end
+
+      def begin_analyze_multiplex(multiplex, analyzers)
+        @nr_analyze = NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: "GraphQL/analyze", category: :web)
         super
       end
 
-      def execute_query(query:)
-        set_this_txn_name =  query.context[:set_new_relic_transaction_name]
-        if set_this_txn_name == true || (set_this_txn_name.nil? && @set_transaction_name)
+      def end_analyze_multiplex(multiplex, analyzers)
+        @nr_analyze.finish
+        super
+      end
+
+      def begin_execute_multiplex(multiplex)
+        query = multiplex.queries.first
+        set_this_txn_name = query.context[:set_new_relic_transaction_name]
+        if set_this_txn_name || (set_this_txn_name.nil? && @set_transaction_name)
           NewRelic::Agent.set_transaction_name(transaction_name(query))
         end
-        NewRelic::Agent::MethodTracerHelpers.trace_execution_scoped("GraphQL/execute") do
-          super
-        end
+        @nr_execute = NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: "GraphQL/execute", category: :web)
+        super
       end
 
-      {
-        "lex" => "GraphQL/lex",
-        "parse" => "GraphQL/parse",
-        "validate" => "GraphQL/validate",
-        "analyze_query" => "GraphQL/analyze",
-        "analyze_multiplex" => "GraphQL/analyze",
-        "execute_multiplex" => "GraphQL/execute",
-        "execute_query_lazy" => "GraphQL/execute",
-      }.each do |trace_method, platform_key|
-        module_eval <<-RUBY, __FILE__, __LINE__
-          def #{trace_method}(**_keys)
-            NewRelic::Agent::MethodTracerHelpers.trace_execution_scoped("#{platform_key}") do
-              super
-            end
-          end
-        RUBY
-      end
-
-      def platform_execute_field(platform_key)
-        NewRelic::Agent::MethodTracerHelpers.trace_execution_scoped(platform_key) do
-          yield
-        end
+      def end_execute_multiplex(multiplex)
+        @nr_execute.finish
+        super
       end
 
-      def platform_authorized(platform_key)
-        NewRelic::Agent::MethodTracerHelpers.trace_execution_scoped(platform_key) do
-          yield
-        end
+      def begin_execute_field(field, object, arguments, query)
+        nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: @nr_field_names[field], category: :web)
+        super
       end
 
-      def platform_resolve_type(platform_key)
-        NewRelic::Agent::MethodTracerHelpers.trace_execution_scoped(platform_key) do
-          yield
-        end
+      def end_execute_field(field, objects, arguments, query, result)
+        nr_segment_stack.pop.finish
+        super
       end
 
-      def platform_field_key(field)
-        "GraphQL/#{field.owner.graphql_name}/#{field.graphql_name}"
+      def begin_authorized(type, obj, ctx)
+        nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: @nr_authorized_names[type], category: :web)
+        super
       end
 
-      def platform_authorized_key(type)
-        "GraphQL/Authorize/#{type.graphql_name}"
+      def end_authorized(type, obj, ctx, is_authed)
+        nr_segment_stack.pop.finish
+        super
+      end
+
+      def begin_resolve_type(type, value, context)
+        nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: @nr_resolve_type_names[type], category: :web)
+        super
+      end
+
+      def end_resolve_type(type, value, context, resolved_type)
+        nr_segment_stack.pop.finish
+        super
+      end
+
+      def begin_dataloader(dl)
+        # nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: "GraphQL/dataloader", category: :web)
+        super
+      end
+
+      def end_dataloader(dl)
+        # nr_segment_stack.pop.finish
+        super
+      end
+
+      def begin_dataloader_source(source)
+        nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: @nr_source_names[source], category: :web)
+        super
+      end
+
+      def end_dataloader_source(source)
+        nr_segment_stack.pop.finish
+        super
+      end
+
+      def dataloader_fiber_yield(source)
+        current_segment = nr_segment_stack.last
+        current_segment.finish
+        super
+      end
+
+      def dataloader_fiber_resume(source)
+        prev_segment = nr_segment_stack.pop
+        seg_partial_name = prev_segment.name.sub(/^.*(GraphQL.*)$/, '\1')
+        nr_segment_stack << NewRelic::Agent::Tracer.start_transaction_or_segment(partial_name: seg_partial_name, category: :web)
+        super
+      end
+
+      private
+
+      def nr_segment_stack
+        Fiber[:graphql_nr_segment_stack] ||= []
+      end
+
+      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 platform_resolve_type_key(type)
-        "GraphQL/ResolveType/#{type.graphql_name}"
+      def fallback_transaction_name(context)
+        context[:tracing_fallback_transaction_name]
       end
     end
   end

lib/graphql/tracing/perfetto_trace.rb

@@ -66,7 +66,6 @@ def initialize(active_support_notifications_pattern: nil, **_rest)
         @sequence_id = object_id
         @pid = Process.pid
         @flow_ids = Hash.new { |h, source_inst| h[source_inst] = [] }.compare_by_identity
-        # TODO intern these too:
         @new_interned_event_names = {}
         @interned_event_name_iids = Hash.new { |h, k|
           new_id = 100 + h.size
@@ -172,7 +171,7 @@ def initialize(active_support_notifications_pattern: nil, **_rest)
         end
       end
 
-      def begin_multiplex(m)
+      def begin_execute_multiplex(m)
         @packets << trace_packet(
           type: TrackEvent::Type::TYPE_SLICE_BEGIN,
           track_uuid: fid,
@@ -184,7 +183,7 @@ def begin_multiplex(m)
         super
       end
 
-      def end_multiplex(m)
+      def end_execute_multiplex(m)
         @packets << trace_packet(
           type: TrackEvent::Type::TYPE_SLICE_END,
           track_uuid: fid,
@@ -286,7 +285,7 @@ def begin_validate(query, validate)
         super
       end
 
-      def end_validate(query, validate, is_valid)
+      def end_validate(query, validate, validation_errors)
         @packets << trace_packet(
           type: TrackEvent::Type::TYPE_SLICE_END,
           track_uuid: fid,
@@ -298,7 +297,7 @@ def end_validate(query, validate, is_valid)
           {
             debug_annotations: [
               @begin_validate.track_event.debug_annotations.first,
-              payload_to_debug("valid?", is_valid)
+              payload_to_debug("valid?", !validation_errors.empty?)
             ]
           }
         )