feat: add support for federated tracing

.rubocop.yml

@@ -45,6 +45,8 @@ Metrics/ParameterLists:
   CountKeywordArgs: false
 Metrics/LineLength:
   Max: 100
+  Exclude:
+    - lib/apollo-federation/tracing/proto/apollo_pb.rb
 
 Naming/FileName:
   Exclude:

Gemfile.lock

@@ -2,6 +2,7 @@ PATH
   remote: .
   specs:
     apollo-federation (0.3.2)
+      google-protobuf
       graphql
 
 GEM
@@ -33,7 +34,8 @@ GEM
     crass (1.0.4)
     diff-lcs (1.3)
     erubi (1.8.0)
-    graphql (1.9.11)
+    google-protobuf (3.9.1)
+    graphql (1.9.12)
     i18n (1.6.0)
       concurrent-ruby (~> 1.0)
     jaro_winkler (1.5.3)

README.md

@@ -147,6 +147,30 @@ class User < BaseObject
 end
 ```
 
+### Tracing
+
+To support [federated tracing](https://www.apollographql.com/docs/apollo-server/federation/metrics/):
+
+1. Add `use ApolloFederation::Tracing` to your schema class.
+2. Change your controller to add `tracing_enabled: true` to the execution context based on the presence of the "include trace" header:
+    ```ruby
+    def execute
+      # ...
+      context = {
+        tracing_enabled: ApolloFederation::Tracing.should_add_traces(headers)
+      }
+      # ...
+    end
+    ```
+3. Change your controller to attach the traces to the response:
+    ```ruby
+    def execute
+      # ...
+      result = YourSchema.execute(query, ...)
+      render json: ApolloFederation::Tracing.attach_trace_to_result(result)
+    end
+    ```
+
 ## Known Issues and Limitations
  - Currently only works with class-based schemas
  - Does not add directives to the output of `Schema.to_definition`. Since `graphql-ruby` doesn't natively support schema directives, the directives will only be visible to the [Apollo Gateway](https://www.apollographql.com/docs/apollo-server/api/apollo-gateway/) through the `Query._service` field (see the [Apollo Federation specification](https://www.apollographql.com/docs/apollo-server/federation/federation-spec/))

apollo-federation.gemspec

@@ -28,6 +28,8 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency 'graphql'
 
+  spec.add_runtime_dependency 'google-protobuf'
+
   spec.add_development_dependency 'actionpack'
   spec.add_development_dependency 'pry-byebug'
   spec.add_development_dependency 'rack'

bin/generate-protos.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+
+set -eo pipefail
+
+DIR=`dirname "$0"`
+OUTPUT_DIR=$DIR/../lib/apollo-federation/tracing/proto
+
+echo "Removing old client"
+rm -f $OUTPUT_DIR/apollo.proto $OUTPUT_DIR/apollo_pb.rb
+
+echo "Downloading latest Apollo Protobuf IDL"
+curl --silent --output lib/apollo-federation/tracing/proto/apollo.proto https://raw.githubusercontent.com/apollographql/apollo-server/master/packages/apollo-engine-reporting-protobuf/src/reports.proto
+
+echo "Generating Ruby client stubs"
+protoc -I lib/apollo-federation/tracing/proto --ruby_out lib/apollo-federation/tracing/proto lib/apollo-federation/tracing/proto/apollo.proto

lib/apollo-federation.rb

@@ -4,3 +4,7 @@
 require 'apollo-federation/schema'
 require 'apollo-federation/object'
 require 'apollo-federation/field'
+require 'apollo-federation/tracing/proto'
+require 'apollo-federation/tracing/node_map'
+require 'apollo-federation/tracing/tracer'
+require 'apollo-federation/tracing'

lib/apollo-federation/tracing.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module ApolloFederation
+  module Tracing
+    KEY = :ftv1
+    DEBUG_KEY = "#{KEY}_debug".to_sym
+
+    module_function
+
+    def use(schema)
+      schema.tracer ApolloFederation::Tracing::Tracer
+    end
+
+    def should_add_traces(headers)
+      headers && headers['apollo-federation-include-trace'] == KEY.to_s
+    end
+
+    def attach_trace_to_result(result)
+      return result unless result.context[:tracing_enabled]
+
+      trace = result.context.namespace(KEY)
+      unless trace[:start_time]
+        raise StandardError.new, 'Apollo Federation Tracing not installed. \
+ Add `use ApollFederation::Tracing` to your schema.'
+      end
+
+      result['errors']&.each do |error|
+        trace[:node_map].add_error(error)
+      end
+
+      proto = ApolloFederation::Tracing::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,
+      )
+
+      result[:extensions] ||= {}
+      result[:extensions][KEY] = Base64.encode64(proto.class.encode(proto))
+
+      if result.context[:debug_tracing]
+        result[:extensions][DEBUG_KEY] = proto.to_h
+      end
+
+      result.to_h
+    end
+
+    def to_proto_timestamp(time)
+      Google::Protobuf::Timestamp.new(seconds: time.to_i, nanos: time.nsec)
+    end
+  end
+end

lib/apollo-federation/tracing/node_map.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require 'active_support/core_ext/array/wrap'
+require 'apollo-federation/tracing/proto'
+
+module ApolloFederation
+  module Tracing
+    # NodeMap stores a flat map of trace nodes by stringified paths
+    # (i.e. "_entities.0.id") for fast lookup when we need to alter
+    # nodes (to add end times or errors.)
+    #
+    # When adding a node to the NodeMap, it will create any missing
+    # parent nodes and ensure the tree is consistent.
+    #
+    # Only the "root" node is attached to the trace extension.
+    class NodeMap
+      ROOT_KEY = ''
+
+      attr_reader :nodes
+      def initialize
+        @nodes = {
+          ROOT_KEY => ApolloFederation::Tracing::Node.new,
+        }
+      end
+
+      def root
+        nodes[ROOT_KEY]
+      end
+
+      def node_for_path(path)
+        nodes[Array.wrap(path).join('.')]
+      end
+
+      def add(path)
+        node = ApolloFederation::Tracing::Node.new
+        node_key = path.join('.')
+        key = path.last
+
+        case key
+        when String # field
+          node.response_name = key
+        when Integer # index of an array
+          node.index = key
+        end
+
+        nodes[node_key] = node
+
+        # find or create a parent node and add this node to its children
+        parent_path = path[0..-2]
+        parent_node = nodes[parent_path.join('.')] || add(parent_path)
+        parent_node.child << node
+
+        node
+      end
+
+      def add_error(error)
+        path = Array.wrap(error['path']).join('.')
+        node = nodes[path] || root
+
+        locations = Array.wrap(error['locations']).map do |location|
+          ApolloFederation::Tracing::Location.new(location)
+        end
+
+        node.error << ApolloFederation::Tracing::Error.new(
+          message: error['message'],
+          location: locations,
+          json: JSON.dump(error),
+        )
+      end
+    end
+  end
+end

lib/apollo-federation/tracing/proto.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require_relative 'proto/apollo_pb'
+
+module ApolloFederation
+  module Tracing
+    Trace = ::Mdg::Engine::Proto::Trace
+    Node = ::Mdg::Engine::Proto::Trace::Node
+    Location = ::Mdg::Engine::Proto::Trace::Location
+    Error = ::Mdg::Engine::Proto::Trace::Error
+  end
+end