feat: add support for federated tracing

Federated GraphQL services should include timing and error
information as a Base64-encoded protocol buffer message in
the `"extensions.ftv1"` field. The gateway requests traces
by adding a special header to the GraphQL request, and combines
traces from all federated services into a single trace.

This change includes a Tracer that uses the graphql-ruby
[tracing API][t] to record field timings and info and store it
on the execution context. It also includes methods on the
`ApolloFederation::Tracing` module to pluck the info from the
context, convert it to an encoded string, and attach it to the
query result's extensions.

I used the Apollo Server typescript code as reference:
* https://github.com/apollographql/apollo-server/blob/master/packages/apollo-engine-reporting/src/federatedExtension.ts
* https://github.com/apollographql/apollo-server/blob/master/packages/apollo-engine-reporting/src/treeBuilder.ts

As well as an unfinished fork of apollo-tracing-ruby:
* https://github.com/salsify/apollo-tracing-ruby/blob/feature/new-apollo-api/lib/apollo_tracing/tracer.rb
* https://github.com/salsify/apollo-tracing-ruby/blob/feature/new-apollo-api/lib/apollo_tracing/trace_tree.rb

Federated tracing documentation: https://www.apollographql.com/docs/apollo-server/federation/metrics/

Addresses Gusto#14

[t]:https://graphql-ruby.org/queries/tracing.html

.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.2.0)
+      google-protobuf
       graphql
 
 GEM
@@ -33,6 +34,7 @@ GEM
     crass (1.0.4)
     diff-lcs (1.3)
     erubi (1.8.0)
+    google-protobuf (3.9.0)
     graphql (1.9.9)
     i18n (1.6.0)
       concurrent-ruby (~> 1.0)

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,50 @@
+# frozen_string_literal: true
+
+module ApolloFederation
+  module Tracing
+    KEY = :ftv1
+
+    def self.use(schema)
+      schema.tracer ApolloFederation::Tracing::Tracer
+    end
+
+    def self.should_add_traces(headers)
+      headers&.['apollo-federation-include-trace'] == KEY.to_s
+    end
+
+    def self.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,
+      )
+
+      json = result.to_h
+      result[:extensions] ||= {}
+      result[:extensions][KEY] = Base64.encode64(proto.class.encode(proto))
+
+      if result.context[:debug_tracing]
+        result[:extensions]["#{KEY}_debug".to_sym] = proto.to_h
+      end
+
+      json
+    end
+
+    def self.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