apollo-studio-tracing 1.0.0.beta.1

data/lib/apollo-studio-tracing/proto/apollo_pb.rb

@@ -0,0 +1,235 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: apollo.proto
+
+require 'google/protobuf'
+
+require 'google/protobuf/timestamp_pb'
+Google::Protobuf::DescriptorPool.generated_pool.build do
+  add_file("apollo.proto", :syntax => :proto3) do
+    add_message "mdg.engine.proto.Trace" do
+      optional :start_time, :message, 4, "google.protobuf.Timestamp"
+      optional :end_time, :message, 3, "google.protobuf.Timestamp"
+      optional :duration_ns, :uint64, 11
+      optional :root, :message, 14, "mdg.engine.proto.Trace.Node"
+      optional :signature, :string, 19
+      optional :details, :message, 6, "mdg.engine.proto.Trace.Details"
+      optional :client_name, :string, 7
+      optional :client_version, :string, 8
+      optional :client_address, :string, 9
+      optional :client_reference_id, :string, 23
+      optional :http, :message, 10, "mdg.engine.proto.Trace.HTTP"
+      optional :cache_policy, :message, 18, "mdg.engine.proto.Trace.CachePolicy"
+      optional :query_plan, :message, 26, "mdg.engine.proto.Trace.QueryPlanNode"
+      optional :full_query_cache_hit, :bool, 20
+      optional :persisted_query_hit, :bool, 21
+      optional :persisted_query_register, :bool, 22
+      optional :registered_operation, :bool, 24
+      optional :forbidden_operation, :bool, 25
+      optional :legacy_signature_needs_resigning, :string, 5
+    end
+    add_message "mdg.engine.proto.Trace.CachePolicy" do
+      optional :scope, :enum, 1, "mdg.engine.proto.Trace.CachePolicy.Scope"
+      optional :max_age_ns, :int64, 2
+    end
+    add_enum "mdg.engine.proto.Trace.CachePolicy.Scope" do
+      value :UNKNOWN, 0
+      value :PUBLIC, 1
+      value :PRIVATE, 2
+    end
+    add_message "mdg.engine.proto.Trace.Details" do
+      map :variables_json, :string, :string, 4
+      map :deprecated_variables, :string, :bytes, 1
+      optional :operation_name, :string, 3
+    end
+    add_message "mdg.engine.proto.Trace.Error" do
+      optional :message, :string, 1
+      repeated :location, :message, 2, "mdg.engine.proto.Trace.Location"
+      optional :time_ns, :uint64, 3
+      optional :json, :string, 4
+    end
+    add_message "mdg.engine.proto.Trace.HTTP" do
+      optional :method, :enum, 1, "mdg.engine.proto.Trace.HTTP.Method"
+      optional :host, :string, 2
+      optional :path, :string, 3
+      map :request_headers, :string, :message, 4, "mdg.engine.proto.Trace.HTTP.Values"
+      map :response_headers, :string, :message, 5, "mdg.engine.proto.Trace.HTTP.Values"
+      optional :status_code, :uint32, 6
+      optional :secure, :bool, 8
+      optional :protocol, :string, 9
+    end
+    add_message "mdg.engine.proto.Trace.HTTP.Values" do
+      repeated :value, :string, 1
+    end
+    add_enum "mdg.engine.proto.Trace.HTTP.Method" do
+      value :UNKNOWN, 0
+      value :OPTIONS, 1
+      value :GET, 2
+      value :HEAD, 3
+      value :POST, 4
+      value :PUT, 5
+      value :DELETE, 6
+      value :TRACE, 7
+      value :CONNECT, 8
+      value :PATCH, 9
+    end
+    add_message "mdg.engine.proto.Trace.Location" do
+      optional :line, :uint32, 1
+      optional :column, :uint32, 2
+    end
+    add_message "mdg.engine.proto.Trace.Node" do
+      optional :original_field_name, :string, 14
+      optional :type, :string, 3
+      optional :parent_type, :string, 13
+      optional :cache_policy, :message, 5, "mdg.engine.proto.Trace.CachePolicy"
+      optional :start_time, :uint64, 8
+      optional :end_time, :uint64, 9
+      repeated :error, :message, 11, "mdg.engine.proto.Trace.Error"
+      repeated :child, :message, 12, "mdg.engine.proto.Trace.Node"
+      oneof :id do
+        optional :response_name, :string, 1
+        optional :index, :uint32, 2
+      end
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode" do
+      oneof :node do
+        optional :sequence, :message, 1, "mdg.engine.proto.Trace.QueryPlanNode.SequenceNode"
+        optional :parallel, :message, 2, "mdg.engine.proto.Trace.QueryPlanNode.ParallelNode"
+        optional :fetch, :message, 3, "mdg.engine.proto.Trace.QueryPlanNode.FetchNode"
+        optional :flatten, :message, 4, "mdg.engine.proto.Trace.QueryPlanNode.FlattenNode"
+      end
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode.SequenceNode" do
+      repeated :nodes, :message, 1, "mdg.engine.proto.Trace.QueryPlanNode"
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode.ParallelNode" do
+      repeated :nodes, :message, 1, "mdg.engine.proto.Trace.QueryPlanNode"
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode.FetchNode" do
+      optional :service_name, :string, 1
+      optional :trace_parsing_failed, :bool, 2
+      optional :trace, :message, 3, "mdg.engine.proto.Trace"
+      optional :sent_time_offset, :uint64, 4
+      optional :sent_time, :message, 5, "google.protobuf.Timestamp"
+      optional :received_time, :message, 6, "google.protobuf.Timestamp"
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode.FlattenNode" do
+      repeated :response_path, :message, 1, "mdg.engine.proto.Trace.QueryPlanNode.ResponsePathElement"
+      optional :node, :message, 2, "mdg.engine.proto.Trace.QueryPlanNode"
+    end
+    add_message "mdg.engine.proto.Trace.QueryPlanNode.ResponsePathElement" do
+      oneof :id do
+        optional :field_name, :string, 1
+        optional :index, :uint32, 2
+      end
+    end
+    add_message "mdg.engine.proto.ReportHeader" do
+      optional :hostname, :string, 5
+      optional :agent_version, :string, 6
+      optional :service_version, :string, 7
+      optional :runtime_version, :string, 8
+      optional :uname, :string, 9
+      optional :schema_tag, :string, 10
+      optional :executable_schema_id, :string, 11
+    end
+    add_message "mdg.engine.proto.PathErrorStats" do
+      map :children, :string, :message, 1, "mdg.engine.proto.PathErrorStats"
+      optional :errors_count, :uint64, 4
+      optional :requests_with_errors_count, :uint64, 5
+    end
+    add_message "mdg.engine.proto.QueryLatencyStats" do
+      repeated :latency_count, :int64, 1
+      optional :request_count, :uint64, 2
+      optional :cache_hits, :uint64, 3
+      optional :persisted_query_hits, :uint64, 4
+      optional :persisted_query_misses, :uint64, 5
+      repeated :cache_latency_count, :int64, 6
+      optional :root_error_stats, :message, 7, "mdg.engine.proto.PathErrorStats"
+      optional :requests_with_errors_count, :uint64, 8
+      repeated :public_cache_ttl_count, :int64, 9
+      repeated :private_cache_ttl_count, :int64, 10
+      optional :registered_operation_count, :uint64, 11
+      optional :forbidden_operation_count, :uint64, 12
+    end
+    add_message "mdg.engine.proto.StatsContext" do
+      optional :client_reference_id, :string, 1
+      optional :client_name, :string, 2
+      optional :client_version, :string, 3
+    end
+    add_message "mdg.engine.proto.ContextualizedQueryLatencyStats" do
+      optional :query_latency_stats, :message, 1, "mdg.engine.proto.QueryLatencyStats"
+      optional :context, :message, 2, "mdg.engine.proto.StatsContext"
+    end
+    add_message "mdg.engine.proto.ContextualizedTypeStats" do
+      optional :context, :message, 1, "mdg.engine.proto.StatsContext"
+      map :per_type_stat, :string, :message, 2, "mdg.engine.proto.TypeStat"
+    end
+    add_message "mdg.engine.proto.FieldStat" do
+      optional :return_type, :string, 3
+      optional :errors_count, :uint64, 4
+      optional :count, :uint64, 5
+      optional :requests_with_errors_count, :uint64, 6
+      repeated :latency_count, :int64, 8
+    end
+    add_message "mdg.engine.proto.TypeStat" do
+      map :per_field_stat, :string, :message, 3, "mdg.engine.proto.FieldStat"
+    end
+    add_message "mdg.engine.proto.Field" do
+      optional :name, :string, 2
+      optional :return_type, :string, 3
+    end
+    add_message "mdg.engine.proto.Type" do
+      optional :name, :string, 1
+      repeated :field, :message, 2, "mdg.engine.proto.Field"
+    end
+    add_message "mdg.engine.proto.Report" do
+      optional :header, :message, 1, "mdg.engine.proto.ReportHeader"
+      map :traces_per_query, :string, :message, 5, "mdg.engine.proto.TracesAndStats"
+      optional :end_time, :message, 2, "google.protobuf.Timestamp"
+    end
+    add_message "mdg.engine.proto.ContextualizedStats" do
+      optional :context, :message, 1, "mdg.engine.proto.StatsContext"
+      optional :query_latency_stats, :message, 2, "mdg.engine.proto.QueryLatencyStats"
+      map :per_type_stat, :string, :message, 3, "mdg.engine.proto.TypeStat"
+    end
+    add_message "mdg.engine.proto.TracesAndStats" do
+      repeated :trace, :message, 1, "mdg.engine.proto.Trace"
+      repeated :stats_with_context, :message, 2, "mdg.engine.proto.ContextualizedStats"
+    end
+  end
+end
+
+module Mdg
+  module Engine
+    module Proto
+      Trace = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace").msgclass
+      Trace::CachePolicy = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.CachePolicy").msgclass
+      Trace::CachePolicy::Scope = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.CachePolicy.Scope").enummodule
+      Trace::Details = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.Details").msgclass
+      Trace::Error = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.Error").msgclass
+      Trace::HTTP = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.HTTP").msgclass
+      Trace::HTTP::Values = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.HTTP.Values").msgclass
+      Trace::HTTP::Method = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.HTTP.Method").enummodule
+      Trace::Location = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.Location").msgclass
+      Trace::Node = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.Node").msgclass
+      Trace::QueryPlanNode = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode").msgclass
+      Trace::QueryPlanNode::SequenceNode = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode.SequenceNode").msgclass
+      Trace::QueryPlanNode::ParallelNode = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode.ParallelNode").msgclass
+      Trace::QueryPlanNode::FetchNode = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode.FetchNode").msgclass
+      Trace::QueryPlanNode::FlattenNode = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode.FlattenNode").msgclass
+      Trace::QueryPlanNode::ResponsePathElement = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Trace.QueryPlanNode.ResponsePathElement").msgclass
+      ReportHeader = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.ReportHeader").msgclass
+      PathErrorStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.PathErrorStats").msgclass
+      QueryLatencyStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.QueryLatencyStats").msgclass
+      StatsContext = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.StatsContext").msgclass
+      ContextualizedQueryLatencyStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.ContextualizedQueryLatencyStats").msgclass
+      ContextualizedTypeStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.ContextualizedTypeStats").msgclass
+      FieldStat = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.FieldStat").msgclass
+      TypeStat = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.TypeStat").msgclass
+      Field = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Field").msgclass
+      Type = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Type").msgclass
+      Report = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.Report").msgclass
+      ContextualizedStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.ContextualizedStats").msgclass
+      TracesAndStats = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("mdg.engine.proto.TracesAndStats").msgclass
+    end
+  end
+end

data/lib/apollo-studio-tracing/shutdown_barrier.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require 'concurrent-ruby'
+
+module ApolloStudioTracing
+  class ShutdownBarrier
+    def initialize
+      @latch = Concurrent::CountDownLatch.new(1)
+    end
+
+    def await_shutdown(timeout_secs)
+      @latch.wait(timeout_secs)
+    end
+
+    def shutdown
+      @latch.count_down
+    end
+
+    def shutting_down?
+      @latch.count.zero?
+    end
+  end
+end

data/lib/apollo-studio-tracing/trace_channel.rb

@@ -0,0 +1,182 @@
+# frozen_string_literal: true
+
+require_relative 'api'
+require_relative 'proto'
+require_relative 'shutdown_barrier'
+
+module ApolloStudioTracing
+  # rubocop:disable Metrics/ClassLength
+  class TraceChannel
+    attr_reader :compress,
+                :api_key,
+                :reporting_interval,
+                :max_uncompressed_report_size,
+                :debug_reports,
+                :max_upload_attempts,
+                :min_upload_retry_delay_secs,
+                :max_queue_bytes
+
+    alias debug_reports? debug_reports
+
+    # rubocop:disable Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
+    def initialize(report_header:, compress: nil, api_key: nil, reporting_interval: nil,
+                   max_uncompressed_report_size: nil, max_queue_bytes: nil, debug_reports: nil,
+                   max_upload_attempts: nil, min_upload_retry_delay_secs: nil)
+      @report_header = report_header
+      @compress = compress.nil? ? true : compress
+      @api_key = api_key || ENV.fetch('ENGINE_API_KEY', ENV.fetch('APOLLO_KEY', 'NO_API_KEY'))
+      @reporting_interval = reporting_interval || 5
+      @max_uncompressed_report_size = max_uncompressed_report_size || 4 * 1024 * 1024
+      @max_queue_bytes = max_queue_bytes || @max_uncompressed_report_size * 10
+      @max_upload_attempts = max_upload_attempts || 5
+      @min_upload_retry_delay_secs = min_upload_retry_delay_secs || 0.1
+      @debug_reports = debug_reports.nil? ? false : debug_reports
+      @queue = Queue.new
+      @queue_bytes = Concurrent::AtomicFixnum.new(0)
+      @queue_full = false
+      @enqueue_mutex = Mutex.new
+      @shutdown_barrier = ApolloStudioTracing::ShutdownBarrier.new
+    end
+    # rubocop:enable Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
+
+    def queue(query_key, trace, context)
+      @enqueue_mutex.synchronize do
+        if @queue_bytes.value >= max_queue_bytes
+          unless @queue_full
+            ApolloStudioTracing.logger.warn(
+              "Apollo tracing queue is above the threshold of #{max_queue_bytes} bytes and " \
+              'trace collection will be paused.',
+            )
+            @queue_full = true
+          end
+        else
+          if @queue_full
+            ApolloStudioTracing.logger.info(
+              "Apollo tracing queue is below the threshold of #{max_queue_bytes} bytes and " \
+              'trace collection will resume.',
+            )
+            @queue_full = false
+          end
+
+          if debug_reports?
+            ApolloStudioTracing.logger.info("Queueing a trace for #{query_key}")
+          end
+
+          proto = ApolloStudioTracing::Trace.new(
+            start_time: to_proto_timestamp(trace[:start_time]),
+            end_time: to_proto_timestamp(trace[:end_time]),
+            duration_ns: trace[:end_time_nanos] - trace[:start_time_nanos],
+            root: trace[:node_map].root,
+            client_name: context[:apollo_client_name],
+            client_version: context[:apollo_client_version],
+          )
+
+          encoded_trace = ApolloStudioTracing::Trace.encode(proto)
+          @queue << [query_key, encoded_trace]
+          @queue_bytes.increment(encoded_trace.bytesize + query_key.bytesize)
+
+          ensure_thread_started
+        end
+      end
+    end
+
+    def start
+      @uploader_thread = Thread.new do
+        run_uploader
+      end
+    end
+
+    def ensure_thread_started
+      return if @uploader_thread&.alive?
+
+      start
+    end
+
+    def flush
+      until @queue.empty?
+        # If the uploader thread isn't running then the queue will never drain
+        break unless @uploader_thread&.alive?
+
+        sleep(0.1)
+      end
+    end
+
+    def shutdown
+      return unless @uploader_thread
+
+      @shutdown_barrier.shutdown
+      @uploader_thread.join
+    end
+
+    private
+
+    def queue_full?
+      @queue_bytes.value >= max_queue_bytes
+    end
+
+    def run_uploader
+      ApolloStudioTracing.logger.info('Apollo trace uploader starting')
+      drain_queue until @shutdown_barrier.await_shutdown(reporting_interval)
+      puts 'Stopping uploader run loop'
+      drain_queue
+    rescue StandardError => e
+      ApolloStudioTracing.logger.warn("Exception thrown in uploader process. #{e}")
+      raise e
+    ensure
+      ApolloStudioTracing.logger.info('Apollo trace uploader exiting')
+    end
+
+    def drain_queue
+      traces_per_query = {}
+      report_size = 0
+      until @queue.empty?
+        query_key, encoded_trace = @queue.pop(false)
+        @queue_bytes.decrement(encoded_trace.bytesize + query_key.bytesize)
+
+        traces_per_query[query_key] ||= []
+        traces_per_query[query_key] << encoded_trace
+        report_size += encoded_trace.bytesize + query_key.bytesize
+
+        next unless report_size >= max_uncompressed_report_size
+
+        send_report(traces_per_query)
+        traces_per_query = {}
+        report_size = 0
+      end
+
+      send_report(traces_per_query) unless traces_per_query.empty?
+    end
+
+    def send_report(traces_per_query)
+      trace_report = ApolloStudioTracing::Report.new(header: @report_header)
+      traces_per_query.each do |query_key, encoded_traces|
+        trace_report.traces_per_query[query_key] = ApolloStudioTracing::TracesAndStats.new(
+          # TODO: Figure out how to use the already encoded traces like Apollo
+          # https://github.com/apollographql/apollo-server/blob/master/packages/apollo-engine-reporting-protobuf/src/index.js
+          trace: encoded_traces.map do |encoded_trace|
+            ApolloStudioTracing::Trace.decode(encoded_trace)
+          end,
+        )
+      end
+
+      if debug_reports?
+        ApolloStudioTracing.logger.info(
+          "Sending trace report:\n#{JSON.pretty_generate(JSON.parse(trace_report.to_json))}",
+        )
+      end
+
+      ApolloStudioTracing::API.upload(
+        ApolloStudioTracing::Report.encode(trace_report),
+        api_key: api_key,
+        compress: compress,
+        max_attempts: max_upload_attempts,
+        min_retry_delay_secs: min_upload_retry_delay_secs,
+      )
+    end
+
+    def to_proto_timestamp(time)
+      Google::Protobuf::Timestamp.new(seconds: time.to_i, nanos: time.nsec)
+    end
+  end
+  # rubocop:enable Metrics/ClassLength
+end

data/lib/apollo-studio-tracing/tracer.rb

@@ -0,0 +1,292 @@
+# frozen_string_literal: true
+
+require 'apollo-studio-tracing/version'
+require 'apollo-studio-tracing/trace_channel'
+
+# Trace events are nested and fire in this order
+# for a simple single-field query like `{ foo }`:
+#
+# <execute_multiplex>
+#   <lex></lex>
+#   <parse></parse>
+#   <validate></validate>
+#   <analyze_multiplex>
+#     <analyze_query></analyze_query>
+#   </analyze_multiplex>
+#
+#   <execute_query>
+#     <execute_field></execute_field>
+#   </execute_query>
+#
+#   <execute_query_lazy>
+#
+#     # `execute_field_lazy` fires *only* when the field is lazy
+#     # (https://graphql-ruby.org/schema/lazy_execution.html)
+#     # so if it fires we should overwrite the ending times recorded
+#     # in `execute_field` to capture the total execution time.
+#
+#     <execute_field_lazy></execute_field_lazy>
+#
+#   </execute_query_lazy>
+#
+#   # `execute_query_lazy` *always* fires, so it's a
+#   # safe place to capture ending times of the full query.
+#
+# </execute_multiplex>
+
+module ApolloStudioTracing
+  # rubocop:disable Metrics/ClassLength
+  class Tracer
+    # store string constants to avoid creating new strings for each call to .trace
+    EXECUTE_MULTIPLEX = 'execute_multiplex'
+    EXECUTE_QUERY = 'execute_query'
+    EXECUTE_QUERY_LAZY = 'execute_query_lazy'
+    EXECUTE_FIELD = 'execute_field'
+    EXECUTE_FIELD_LAZY = 'execute_field_lazy'
+
+    attr_reader :trace_prepare, :query_signature
+
+    def initialize(
+      schema_tag: nil,
+      executable_schema_id: nil,
+      service_version: nil,
+      trace_prepare: nil,
+      query_signature: nil,
+      api_key: nil,
+      **trace_channel_options
+    )
+      @trace_prepare = trace_prepare || proc {}
+      @query_signature = query_signature || proc do |query|
+        # TODO: This should be smarter
+        # TODO (lsanwick) Replace with reference implementation from
+        # https://github.com/apollographql/apollo-tooling/blob/master/packages/apollo-graphql/src/operationId.ts
+        query.query_string
+      end
+
+      report_header = ApolloStudioTracing::ReportHeader.new(
+        hostname: hostname,
+        agent_version: agent_version,
+        service_version: service_version,
+        runtime_version: RUBY_DESCRIPTION,
+        uname: uname,
+        schema_tag: schema_tag || ENV.fetch('ENGINE_SCHEMA_TAG', 'current'),
+        executable_schema_id: executable_schema_id,
+      )
+      @trace_channel = ApolloStudioTracing::TraceChannel.new(
+        report_header: report_header,
+        api_key: api_key,
+        **trace_channel_options,
+      )
+    end
+
+    def start_trace_channel
+      @trace_channel.start
+    end
+
+    def shutdown_trace_channel
+      @trace_channel.shutdown
+    end
+
+    def flush_trace_channel
+      @trace_channel.flush
+    end
+
+    def trace(key, data, &block)
+      case key
+      when EXECUTE_MULTIPLEX
+        execute_multiplex(data, &block)
+      when EXECUTE_QUERY_LAZY
+        execute_query_lazy(data, &block)
+      when EXECUTE_FIELD
+        execute_field(data, &block)
+      when EXECUTE_FIELD_LAZY
+        execute_field_lazy(data, &block)
+      else
+        yield
+      end
+    end
+
+    def tracing_enabled?(context)
+      context && context[:apollo_tracing_enabled]
+    end
+
+    def execute_multiplex(data, &block)
+      # Step 1:
+      # Create a trace hash on each query's context and record start times.
+      data.fetch(:multiplex).queries.each { |query| start_trace(query) }
+
+      results = block.call
+
+      # Step 5
+      #  Enqueue the final trace onto the TraceChannel.
+      results.map { |result| attach_trace_to_result(result) }
+    end
+
+    def start_trace(query)
+      return unless tracing_enabled?(query&.context)
+
+      query.context.namespace(ApolloStudioTracing::KEY).merge!(
+        start_time: Time.now.utc,
+        start_time_nanos: Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond),
+        node_map: NodeMap.new,
+      )
+    end
+
+    # Step 2:
+    # * Record start and end times for the field resolver.
+    # * Rescue errors so the method doesn't exit early.
+    # * Create a trace "node" and attach field details.
+    # * Propagate the error (if necessary) so it ends up in the top-level errors array.
+    #
+    # The values in `data` are different depending on the executor runtime.
+    # https://graphql-ruby.org/api-doc/1.9.3/GraphQL/Tracing
+    #
+    # Nodes are added the NodeMap stored in the trace hash.
+    #
+    # Errors are added to nodes in `ApolloStudioTracing::Tracing.attach_trace_to_result`
+    # because we don't have the error `location` here.
+    # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+    def execute_field(data, &block)
+      context = data.fetch(:context, nil) || data.fetch(:query).context
+      return block.call unless tracing_enabled?(context)
+
+      start_time_nanos = Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
+
+      begin
+        result = block.call
+      rescue StandardError => e
+        error = e
+      end
+
+      end_time_nanos = Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
+
+      # legacy runtime
+      if data.include?(:context)
+        path = context.path
+        field_name = context.field.graphql_name
+        field_type = context.field.type.to_s
+        parent_type = context.parent_type.graphql_name
+      else # interpreter runtime
+        path = data.fetch(:path)
+        field = data.fetch(:field)
+        field_name = field.graphql_name
+        field_type = field.type.to_type_signature
+        parent_type = data.fetch(:owner).graphql_name
+      end
+
+      trace = context.namespace(ApolloStudioTracing::KEY)
+      node = trace[:node_map].add(path)
+
+      # original_field_name is set only for aliased fields
+      node.original_field_name = field_name if field_name != path.last
+      node.type = field_type
+      node.parent_type = parent_type
+      node.start_time = start_time_nanos - trace[:start_time_nanos]
+      node.end_time = end_time_nanos - trace[:start_time_nanos]
+
+      raise error if error
+
+      result
+    end
+
+    # Optional Step 3:
+    # Overwrite the end times on the trace node if the resolver was lazy.
+    def execute_field_lazy(data, &block)
+      context = data.fetch(:context, nil) || data.fetch(:query).context
+      return block.call unless tracing_enabled?(context)
+
+      begin
+        result = block.call
+      rescue StandardError => e
+        error = e
+      end
+
+      end_time_nanos = Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
+
+      # legacy runtime
+      if data.include?(:context)
+        path = context.path
+        field = context.field
+      else # interpreter runtime
+        path = data.fetch(:path)
+        field = data.fetch(:field)
+      end
+
+      trace = context.namespace(ApolloStudioTracing::KEY)
+
+      # When a field is resolved with an array of lazy values, the interpreter fires an
+      # `execute_field` for the resolution of the field and then a `execute_field_lazy` event for
+      # each lazy value in the array. Since the path here will contain an index (indicating which
+      # lazy value we're executing: e.g. ['arrayOfLazies', 0]), we won't have a node for the path.
+      # We only care about the end of the parent field (e.g. ['arrayOfLazies']), so we get the
+      # node for that path. What ends up happening is we update the end_time for the parent node
+      # for each of the lazy values. The last one that's executed becomes the final end time.
+      if field.type.list? && path.last.is_a?(Integer)
+        path = path[0...-1]
+      end
+      node = trace[:node_map].node_for_path(path)
+      node.end_time = end_time_nanos - trace[:start_time_nanos]
+
+      raise error if error
+
+      result
+    end
+    # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+
+    # Step 4:
+    # Record end times and merge them into the trace hash
+    def execute_query_lazy(data, &block)
+      result = block.call
+
+      # Normalize to an array of queries regardless of whether we are multiplexing or performing a
+      # single query.
+      queries = Array(data.fetch(:multiplex)&.queries || data.fetch(:query))
+
+      queries.map do |query|
+        next unless tracing_enabled?(query&.context)
+
+        trace = query.context.namespace(ApolloStudioTracing::KEY)
+
+        trace.merge!(
+          end_time: Time.now.utc,
+          end_time_nanos: Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond),
+        )
+      end
+
+      result
+    end
+
+    private
+
+    def attach_trace_to_result(result)
+      return result unless tracing_enabled?(result.context)
+
+      trace = result.context.namespace(ApolloStudioTracing::KEY)
+
+      result['errors']&.each do |error|
+        trace[:node_map].add_error(error)
+      end
+
+      @trace_channel.queue(
+        "# #{result.query.operation_name || '-'}\n#{query_signature.call(result.query)}",
+        trace,
+        result.context,
+      )
+
+      result
+    end
+
+    def hostname
+      @hostname ||= Socket.gethostname
+    end
+
+    def agent_version
+      @agent_version ||= "apollo-studio-tracing #{ApolloStudioTracing::VERSION}"
+    end
+
+    def uname
+      @uname ||= `uname -a`
+    end
+  end
+  # rubocop:enable Metrics/ClassLength
+end