Merge pull request #5240 from rmosolgo/better-nr-trace

Rewrite NewRelicTrace to support fiber stops/starts

Author: rmosolgo

.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

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+require 'rubocop'
+
+module Cop
+  module Development
+    class TraceMethodsCop < RuboCop::Cop::Base
+      extend RuboCop::Cop::AutoCorrector
+
+      TRACE_HOOKS = [
+        :analyze_multiplex,
+        :analyze_query,
+        :authorized,
+        :authorized_lazy,
+        :begin_analyze_multiplex,
+        :begin_authorized,
+        :begin_dataloader,
+        :begin_dataloader_source,
+        :begin_execute_field,
+        :begin_execute_multiplex,
+        :begin_parse,
+        :begin_resolve_type,
+        :begin_validate,
+        :dataloader_fiber_exit,
+        :dataloader_fiber_resume,
+        :dataloader_fiber_yield,
+        :dataloader_spawn_execution_fiber,
+        :dataloader_spawn_source_fiber,
+        :end_analyze_multiplex,
+        :end_authorized,
+        :end_dataloader,
+        :end_dataloader_source,
+        :end_execute_field,
+        :end_execute_multiplex,
+        :end_parse,
+        :end_resolve_type,
+        :end_validate,
+        :execute_field,
+        :execute_field_lazy,
+        :execute_multiplex,
+        :execute_query,
+        :execute_query_lazy,
+        :lex,
+        :parse,
+        :resolve_type,
+        :resolve_type_lazy,
+        :validate,
+      ]
+
+      MSG = "Trace methods should call `super` to pass control to other traces"
+
+      def on_def(node)
+        if TRACE_HOOKS.include?(node.method_name) && !node.each_descendant(:super, :zsuper).any?
+          add_offense(node) do |corrector|
+            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 = [
+            # Not really necessary for making a good trace:
+            :lex, :analyze_query, :execute_query, :execute_query_lazy,
+            # Only useful for isolated event tracking:
+            :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
+          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
+    end
+  end
+end

cop/development/trace_methods_cop.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/execution/interpreter.rb

@@ -789,8 +789,10 @@ def after_lazy(lazy_obj, field:, owner_object:, arguments:, ast_node:, result:,
               runtime_state.was_authorized_by_scope_items = was_authorized_by_scope_items
               # Wrap the execution of _this_ method with tracing,
               # but don't wrap the continuation below
+              result = nil
               inner_obj = begin
-                if trace
+                result = if trace
+                  @current_trace.begin_execute_field(field, owner_object, arguments, query)
                   @current_trace.execute_field_lazy(field: field, query: query, object: owner_object, arguments: arguments, ast_node: ast_node) do
                     schema.sync_lazy(lazy_obj)
                   end
@@ -805,6 +807,10 @@ def after_lazy(lazy_obj, field:, owner_object:, arguments:, ast_node:, result:,
                 rescue GraphQL::ExecutionError => ex_err
                   ex_err
                 end
+              ensure
+                if trace
+                  @current_trace.end_execute_field(field, owner_object, arguments, query, result)
+                end
               end
               yield(inner_obj, runtime_state)
             end
@@ -852,17 +858,18 @@ def resolve_type(type, value)
           resolved_type, resolved_value = @current_trace.resolve_type(query: query, type: type, object: value) do
             query.resolve_type(type, value)
           end
+          @current_trace.end_resolve_type(type, value, context, resolved_type)
 
           if lazy?(resolved_type)
             GraphQL::Execution::Lazy.new do
+              @current_trace.begin_resolve_type(type, value, context)
               @current_trace.resolve_type_lazy(query: query, type: type, object: value) do
                 rt = schema.sync_lazy(resolved_type)
                 @current_trace.end_resolve_type(type, value, context, rt)
                 rt
               end
             end
           else
-            @current_trace.end_resolve_type(type, value, context, resolved_type)
             [resolved_type, resolved_value]
           end
         end

lib/graphql/execution/interpreter/runtime.rb

@@ -66,29 +66,35 @@ def wrap(object, context)
         # @return [GraphQL::Schema::Object, GraphQL::Execution::Lazy]
         # @raise [GraphQL::UnauthorizedError] if the user-provided hook returns `false`
         def authorized_new(object, context)
-          maybe_lazy_auth_val = context.query.current_trace.authorized(query: context.query, type: self, object: object) do
-            begin
-              context.query.current_trace.begin_authorized(self, object, context)
-              authorized?(object, context)
-            rescue GraphQL::UnauthorizedError => err
-              context.schema.unauthorized_object(err)
-            rescue StandardError => err
-              context.query.handle_or_reraise(err)
+          context.query.current_trace.begin_authorized(self, object, context)
+          begin
+            maybe_lazy_auth_val = context.query.current_trace.authorized(query: context.query, type: self, object: object) do
+              begin
+                authorized?(object, context)
+              rescue GraphQL::UnauthorizedError => err
+                context.schema.unauthorized_object(err)
+              rescue StandardError => err
+                context.query.handle_or_reraise(err)
+              end
             end
+          ensure
+            context.query.current_trace.end_authorized(self, object, context, maybe_lazy_auth_val)
           end
 
           auth_val = if context.schema.lazy?(maybe_lazy_auth_val)
             GraphQL::Execution::Lazy.new do
+              context.query.current_trace.begin_authorized(self, object, context)
               context.query.current_trace.authorized_lazy(query: context.query, type: self, object: object) do
-                context.schema.sync_lazy(maybe_lazy_auth_val)
+                res = context.schema.sync_lazy(maybe_lazy_auth_val)
+                context.query.current_trace.end_authorized(self, object, context, res)
+                res
               end
             end
           else
             maybe_lazy_auth_val
           end
 
           context.query.after_lazy(auth_val) do |is_authorized|
-            context.query.current_trace.end_authorized(self, object, context, is_authorized)
             if is_authorized
               self.new(object, context)
             else

lib/graphql/schema/object.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.