checkpoint towards including traces in federated service responses

maybe works, but entirely untested and lacks the http header to toggle it

packages/apollo-engine-reporting/src/federatedExtension.ts

@@ -0,0 +1,74 @@
+import { GraphQLResolveInfo, GraphQLError } from 'graphql';
+import { GraphQLExtension } from 'graphql-extensions';
+import { Trace } from 'apollo-engine-reporting-protobuf';
+
+import { EngineReportingTreeBuilder } from './treeBuilder';
+
+interface FederatedTraceV1 {
+  d: number;
+  t: string; // base64 encoding of protobuf of Trace.Node
+}
+
+export class EngineFederatedTracingExtension<TContext = any>
+  implements GraphQLExtension<TContext> {
+  public trace = new Trace();
+  private treeBuilder: EngineReportingTreeBuilder;
+  private result?: { durationNs: number; rootNode: Trace.Node };
+
+  public constructor(options: {
+    rewriteError?: (err: GraphQLError) => GraphQLError | null;
+  }) {
+    this.treeBuilder = new EngineReportingTreeBuilder({
+      rewriteError: options.rewriteError,
+    });
+  }
+
+  public requestDidStart(_options: any) {
+    // XXX only do things if we get the header
+    this.treeBuilder.startTiming();
+  }
+
+  public willResolveField(
+    _source: any,
+    _args: { [argName: string]: any },
+    _context: TContext,
+    info: GraphQLResolveInfo,
+  ): ((error: Error | null, result: any) => void) | void {
+    return this.treeBuilder.willResolveField(info);
+  }
+
+  public didEncounterErrors(errors: GraphQLError[]) {
+    this.treeBuilder.didEncounterErrors(errors);
+  }
+
+  public executionDidStart() {
+    // It's a little odd that we record the end time after execution rather than
+    // at the end of the whole request, but because we need to include our
+    // formatted trace in the request itself, we have to record it before the
+    // request is over!  It's also odd that we don't do traces for parse or
+    // validation errors, but runQuery doesn't currently support that, as
+    // format() is only invoked after execution.
+    return () => {
+      this.result = this.treeBuilder.stopTiming();
+    };
+  }
+
+  public format(): [string, FederatedTraceV1] {
+    if (!this.result) {
+      throw Error('format called before end of execution?');
+    }
+    const encodedUint8Array = Trace.Node.encode(this.result.rootNode).finish();
+    const encodedBuffer = Buffer.from(
+      encodedUint8Array,
+      encodedUint8Array.byteOffset,
+      encodedUint8Array.byteLength,
+    );
+    return [
+      'ftv1',
+      {
+        d: this.result.durationNs,
+        t: encodedBuffer.toString('base64'),
+      },
+    ];
+  }
+}

packages/apollo-engine-reporting/src/index.ts

@@ -1 +1,2 @@
 export { EngineReportingOptions, EngineReportingAgent } from './agent';
+export { EngineFederatedTracingExtension } from './federatedExtension';

packages/apollo-server-core/src/ApolloServer.ts

@@ -14,6 +14,8 @@ import {
   ValidationContext,
   FieldDefinitionNode,
   DocumentNode,
+  isObjectType,
+  isScalarType,
 } from 'graphql';
 import { GraphQLExtension } from 'graphql-extensions';
 import {
@@ -363,9 +365,18 @@ export class ApolloServerBase {
         .digest('hex');
     }
 
+    const schemaIsFederated = this.schemaIsFederated();
     if (this.engineServiceId) {
+      if (schemaIsFederated) {
+        // XXX should this throw instead? or be left out?
+        console.warn(
+          "It looks like you're running a federated schema and you've configured your service " +
+            'to report metrics to Apollo Engine. You should only configure your Apollo gateway ' +
+            'to report metrics to Apollo Engine.',
+        );
+      }
       const { EngineReportingAgent } = require('apollo-engine-reporting');
-      this.engineReportingAgent = new EngineReportingAgent(
+      const engineReportingAgent = new EngineReportingAgent(
         typeof engine === 'object' ? engine : Object.create(null),
         {
           schema: this.schema,
@@ -375,8 +386,21 @@ export class ApolloServerBase {
           },
         },
       );
+      this.engineReportingAgent = engineReportingAgent;
       // Let's keep this extension second so it wraps everything, except error formatting
-      this.extensions.push(() => this.engineReportingAgent!.newExtension());
+      this.extensions.push(() => engineReportingAgent.newExtension());
+    } else if (engine !== false && this.schemaIsFederated()) {
+      // We haven't configured this app to use Engine directly. But it looks like
+      // we are a federated service backend, so we should be capable of including
+      // our trace in a response extension if we are asked to by the gateway.
+      const {
+        EngineFederatedTracingExtension,
+      } = require('apollo-engine-reporting');
+      const rewriteError =
+        engine && typeof engine === 'object' ? engine.rewriteError : undefined;
+      this.extensions.push(
+        () => new EngineFederatedTracingExtension({ rewriteError }),
+      );
     }
 
     if (extensions) {
@@ -519,6 +543,31 @@ export class ApolloServerBase {
     return false;
   }
 
+  // Returns true if it appears that the schema was returned from
+  // @apollo/federation's buildFederatedSchema. This strategy avoids depending
+  // explicitly on @apollo/federation or relying on something that might not
+  // survive transformations like monkey-patching a boolean field onto the
+  // schema.
+  //
+  // The only thing this is used for is determining whether traces should be
+  // added to responses if requested with an HTTP header; if there's a false
+  // positive, that feature can be disabled by specifying `engine: false`.
+  private schemaIsFederated(): boolean {
+    const serviceType = this.schema.getType('_Service');
+    if (!(serviceType && isObjectType(serviceType))) {
+      return false;
+    }
+    const sdlField = serviceType.getFields().sdl;
+    if (!sdlField) {
+      return false;
+    }
+    const sdlFieldType = sdlField.type;
+    if (!isScalarType(sdlFieldType)) {
+      return false;
+    }
+    return sdlFieldType.name == 'String';
+  }
+
   private ensurePluginInstantiation(plugins?: PluginDefinition[]): void {
     if (!plugins || !plugins.length) {
       return;