Always send large traces as stats (#6897)

Send large (>10mb) traces as stats always.

Co-authored-by: David Glasser <[email protected]>

Author: bonnici

.changeset/yellow-actors-do.md

@@ -0,0 +1,5 @@
+---
+'@apollo/server': patch
+---
+
+Usage reporting: always send traces over 10MB as stats.

packages/server/src/__tests__/plugin/usageReporting/stats.test.ts

@@ -1,8 +1,13 @@
-import { Trace } from '@apollo/usage-reporting-protobuf';
+import {
+  Trace,
+  ReportHeader,
+  ReferencedFieldsForType,
+} from '@apollo/usage-reporting-protobuf';
 import { dateToProtoTimestamp } from '../../../plugin/traceTreeBuilder';
 import {
   OurContextualizedStats,
   SizeEstimator,
+  OurReport,
 } from '../../../plugin/usageReporting/stats';
 import { DurationHistogram } from '../../../plugin/usageReporting/durationHistogram';
 import { describe, it, expect } from '@jest/globals';
@@ -473,3 +478,81 @@ describe('Check type stats', () => {
     expect(contextualizedStats).toMatchSnapshot();
   });
 });
+
+describe('Add trace to report', () => {
+  const defaultHeader = new ReportHeader({
+    hostname: 'hostname',
+    agentVersion: `@apollo/server`,
+    runtimeVersion: `node latest`,
+    uname: 'uname',
+    executableSchemaId: 'schema',
+    graphRef: 'graph',
+  });
+  const referencedFieldsByType = Object.create(null);
+  referencedFieldsByType['type'] = new ReferencedFieldsForType({
+    fieldNames: ['field1', 'field2'],
+    isInterface: false,
+  });
+
+  it('add as stats if asTrace is false', () => {
+    const report = new OurReport(defaultHeader);
+    report.addTrace({
+      statsReportKey: 'key',
+      trace: baseTrace,
+      asTrace: false,
+      referencedFieldsByType,
+    });
+
+    expect(report.tracesPerQuery['key']?.trace?.length).toBe(0);
+    expect(
+      Object.keys(report.tracesPerQuery['key']?.statsWithContext?.map).length,
+    ).toBe(1);
+  });
+
+  it('add as stats if asTrace is true but trace is too large', () => {
+    const report = new OurReport(defaultHeader);
+    report.addTrace({
+      statsReportKey: 'key',
+      trace: baseTrace,
+      asTrace: true,
+      referencedFieldsByType,
+      maxTraceBytes: 10,
+    });
+
+    expect(report.tracesPerQuery['key']?.trace?.length).toBe(0);
+    expect(
+      Object.keys(report.tracesPerQuery['key']?.statsWithContext?.map).length,
+    ).toBe(1);
+  });
+
+  it('add as trace if asTrace is true and trace is not too large', () => {
+    const report = new OurReport(defaultHeader);
+    report.addTrace({
+      statsReportKey: 'key',
+      trace: baseTrace,
+      asTrace: true,
+      referencedFieldsByType,
+      maxTraceBytes: 500 * 1024,
+    });
+
+    expect(report.tracesPerQuery['key']?.trace?.length).toBe(1);
+    expect(
+      Object.keys(report.tracesPerQuery['key']?.statsWithContext?.map).length,
+    ).toBe(0);
+  });
+
+  it('add as trace if asTrace is true and trace is not too large and max trace size is left as default', () => {
+    const report = new OurReport(defaultHeader);
+    report.addTrace({
+      statsReportKey: 'key',
+      trace: baseTrace,
+      asTrace: true,
+      referencedFieldsByType,
+    });
+
+    expect(report.tracesPerQuery['key']?.trace?.length).toBe(1);
+    expect(
+      Object.keys(report.tracesPerQuery['key']?.statsWithContext?.map).length,
+    ).toBe(0);
+  });
+});

packages/server/src/plugin/usageReporting/defaultSendOperationsAsTrace.ts

@@ -9,7 +9,10 @@ export function defaultSendOperationsAsTrace() {
   // operation, what minute the operation ended at, etc) to `true` if we've seen
   // it recently. We actually split this into one cache per minute so we can
   // throw away a full minute's worth of cache at once; we keep only the last
-  // three minutes
+  // three minutes.
+  // Note that if a trace is over a certain size, we will always send it as
+  // stats. We check this within the addTrace function of the OurReport class so
+  // that we don't have to encode these large traces twice.
   const cache = new LRUCache<string, true>({
     // 3MiB limit, very much approximately since we can't be sure how V8 might
     // be storing these strings internally. Though this should be enough to

packages/server/src/plugin/usageReporting/plugin.ts

@@ -680,7 +680,7 @@ export function ApolloServerPluginUsageReporting<TContext extends BaseContext>(
               trace,
               // We include the operation as a trace (rather than aggregated
               // into stats) only if the user didn't set `sendTraces: false`
-              // *and8 we believe it's possible that our organization's plan
+              // *and* we believe it's possible that our organization's plan
               // allows for viewing traces *and* we actually captured this as
               // a full trace *and* sendOperationAsTrace says so.
               //

packages/server/src/plugin/usageReporting/stats.ts

@@ -61,20 +61,30 @@ export class OurReport implements Required<IReport> {
     trace,
     asTrace,
     referencedFieldsByType,
+    // The max size a trace can be before it is sent as stats. Note that the
+    // Apollo reporting ingress server will never store any traces over 10mb
+    // anyway. They will still be converted to stats as we would do here.
+    maxTraceBytes = 10 * 1024 * 1024,
   }: {
     statsReportKey: string;
     trace: Trace;
     asTrace: boolean;
     referencedFieldsByType: ReferencedFieldsByType;
+    maxTraceBytes?: number;
   }) {
     const tracesAndStats = this.getTracesAndStats({
       statsReportKey,
       referencedFieldsByType,
     });
     if (asTrace) {
       const encodedTrace = Trace.encode(trace).finish();
-      tracesAndStats.trace.push(encodedTrace);
-      this.sizeEstimator.bytes += 2 + encodedTrace.length;
+
+      if (!isNaN(maxTraceBytes) && encodedTrace.length > maxTraceBytes) {
+        tracesAndStats.statsWithContext.addTrace(trace, this.sizeEstimator);
+      } else {
+        tracesAndStats.trace.push(encodedTrace);
+        this.sizeEstimator.bytes += 2 + encodedTrace.length;
+      }
     } else {
       tracesAndStats.statsWithContext.addTrace(trace, this.sizeEstimator);
     }