feat: Enable tracing via OpenTelemetry. (#2218)

* feat: Enable tracing via OpenTelemetry.

* Better comment for the tracerProvider setting.

* Address feedback.

* Add an event for logical termination and add events to tests.

* Address API naming feedback.

* Remove the trace event that's currently not in Firestore prod.

Author: ehsannas

dev/src/index.ts

@@ -792,8 +792,7 @@ export class Firestore implements firestore.Firestore {
   }
 
   private newTraceUtilInstance(settings: firestore.Settings): TraceUtil {
-    // Take the tracing option from the settings.
-    let createEnabledInstance = settings.openTelemetryOptions?.enableTracing;
+    let createEnabledInstance = true;
 
     // The environment variable can override options to enable/disable telemetry collection.
     if ('FIRESTORE_ENABLE_TRACING' in process.env) {

dev/src/reference/query-util.ts

@@ -180,6 +180,7 @@ export class QueryUtil<
     const tag = requestTag();
     const startTime = Date.now();
     const isExplain = explainOptions !== undefined;
+    const methodName = 'runQuery';
 
     let numDocumentsReceived = 0;
     let lastReceivedDocument: QueryDocumentSnapshot<
@@ -245,6 +246,11 @@ export class QueryUtil<
 
         if (proto.done) {
           logger('QueryUtil._stream', tag, 'Trigger Logical Termination.');
+          this._firestore._traceUtil
+            .currentSpan()
+            .addEvent(
+              `Firestore.${methodName}: Received RunQueryResponse.Done.`
+            );
           backendStream.unpipe(stream);
           backendStream.resume();
           backendStream.end();
@@ -265,7 +271,6 @@ export class QueryUtil<
         let streamActive: Deferred<boolean>;
         do {
           streamActive = new Deferred<boolean>();
-          const methodName = 'runQuery';
 
           this._firestore._traceUtil
             .currentSpan()

dev/src/telemetry/disabled-trace-util.ts

@@ -16,6 +16,10 @@
 import {Attributes, TraceUtil} from './trace-util';
 import {Span} from './span';
 
+/**
+ * @private
+ * @internal
+ */
 export class DisabledTraceUtil implements TraceUtil {
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   startSpan(name: string): Span {

dev/src/telemetry/enabled-trace-util.ts

@@ -22,6 +22,7 @@ import {
   trace,
   Tracer,
   Span as OpenTelemetrySpan,
+  TracerProvider,
 } from '@opentelemetry/api';
 
 import {Span} from './span';
@@ -33,22 +34,40 @@ import {DEFAULT_DATABASE_ID} from '../path';
 import {DEFAULT_MAX_IDLE_CHANNELS} from '../index';
 const serviceConfig = interfaces['google.firestore.v1.Firestore'];
 
+/**
+ * @private
+ * @internal
+ */
 export class EnabledTraceUtil implements TraceUtil {
   private tracer: Tracer;
   private settingsAttributes: Attributes;
 
+  // Visible for testing
+  tracerProvider: TracerProvider;
+
   constructor(settings: Settings) {
-    let tracerProvider = settings.openTelemetryOptions?.tracerProvider;
+    let provider: TracerProvider | undefined =
+      settings.openTelemetry?.tracerProvider;
 
     // If a TracerProvider has not been given to us, we try to use the global one.
-    if (!tracerProvider) {
+    if (!provider) {
       const {trace} = require('@opentelemetry/api');
-      tracerProvider = trace.getTracerProvider();
+      provider = trace.getTracerProvider();
     }
 
+    // At this point provider is guaranteed to be defined because
+    // `trace.getTracerProvider()` does not return null or undefined.
+    this.tracerProvider = provider!;
+
     const libVersion = require('../../../package.json').version;
     const libName = require('../../../package.json').name;
-    this.tracer = tracerProvider.getTracer(libName, libVersion);
+    try {
+      this.tracer = this.tracerProvider.getTracer(libName, libVersion);
+    } catch (e) {
+      throw new Error(
+        "The object provided for 'tracerProvider' does not conform to the TracerProvider interface."
+      );
+    }
 
     this.settingsAttributes = {};
     this.settingsAttributes['otel.scope.name'] = libName;

dev/src/telemetry/span.ts

@@ -17,6 +17,10 @@
 import {Span as OpenTelemetrySpan} from '@opentelemetry/api';
 import {Attributes} from './trace-util';
 
+/**
+ * @private
+ * @internal
+ */
 export class Span {
   constructor(private span?: OpenTelemetrySpan) {}
 

dev/src/telemetry/trace-util.ts

@@ -16,9 +16,18 @@
 
 import {Span} from './span';
 
+/**
+ * @private
+ * @internal
+ */
 export interface Attributes {
   [attributeKey: string]: AttributeValue | undefined;
 }
+
+/**
+ * @private
+ * @internal
+ */
 export declare type AttributeValue =
   | string
   | number
@@ -67,6 +76,10 @@ export const ATTRIBUTE_KEY_TRANSACTION_TYPE = 'transaction_type';
 export const ATTRIBUTE_KEY_ATTEMPTS_ALLOWED = 'attempts_allowed';
 export const ATTRIBUTE_KEY_ATTEMPTS_REMAINING = 'attempts_remaining';
 
+/**
+ * @private
+ * @internal
+ */
 export interface TraceUtil {
   startActiveSpan<F extends (span: Span) => unknown>(
     name: string,

dev/system-test/tracing.ts

@@ -30,7 +30,7 @@ import {
   Context as OpenTelemetryContext,
 } from '@opentelemetry/api';
 import {TraceExporter} from '@google-cloud/opentelemetry-cloud-trace-exporter';
-import {Settings} from '@google-cloud/firestore';
+import {FirestoreOpenTelemetryOptions, Settings} from '@google-cloud/firestore';
 import {
   AlwaysOnSampler,
   BatchSpanProcessor,
@@ -96,13 +96,6 @@ setLogFunction((msg: string) => {
   console.log(`LOG: ${msg}`);
 });
 
-// TODO(tracing): This should be moved to firestore.d.ts when we want to
-//  release the feature.
-export interface FirestoreOpenTelemetryOptions {
-  enableTracing?: boolean;
-  tracerProvider?: any;
-}
-
 interface TestConfig {
   // In-Memory tests check trace correctness by inspecting traces in memory by
   // utilizing InMemorySpanExporter. These tests have `e2e` set to `false`.
@@ -192,7 +185,6 @@ describe('Tracing Tests', () => {
     tracerProvider: TracerProvider
   ): FirestoreOpenTelemetryOptions {
     const options: FirestoreOpenTelemetryOptions = {
-      enableTracing: true,
       tracerProvider: undefined,
     };
 
@@ -285,7 +277,7 @@ describe('Tracing Tests', () => {
 
     const settings: Settings = {
       preferRest: testConfig.preferRest,
-      openTelemetryOptions: getOpenTelemetryOptions(tracerProvider),
+      openTelemetry: getOpenTelemetryOptions(tracerProvider),
     };
 
     // Named-database tests use an environment variable to specify the database ID. Add it to the settings.
@@ -660,7 +652,7 @@ describe('Tracing Tests', () => {
 
     // Expect that the span exists first.
     const span = getSpanByName(spanName);
-    expect(span).to.not.be.null;
+    expect(span, `Could not find the span named ${spanName}`).to.not.be.null;
 
     // Assert that the expected attributes are present in the span attributes.
     // Note that the span attributes may be a superset of the attributes passed
@@ -672,6 +664,29 @@ describe('Tracing Tests', () => {
     }
   }
 
+  // Ensures that the given span exists and has the given attributes.
+  function expectSpanHasEvents(spanName: string, eventNames: string[]): void {
+    // The Cloud Trace API does not return span attributes and events.
+    if (testConfig.e2e) {
+      return;
+    }
+
+    // Expect that the span exists first.
+    const span = getSpanByName(spanName);
+    expect(span, `Could not find the span named ${spanName}`).to.not.be.null;
+
+    // Assert that the expected attributes are present in the span attributes.
+    // Note that the span attributes may be a superset of the attributes passed
+    // to this function.
+    if (span?.events) {
+      const numEvents = eventNames.length;
+      expect(numEvents).to.equal(span.events.length);
+      for (let i = 0; i < numEvents; ++i) {
+        expect(span.events[i].name).to.equal(eventNames[i]);
+      }
+    }
+  }
+
   describe(IN_MEMORY_TEST_SUITE_TITLE, () => {
     describe(NON_GLOBAL_OTEL_TEST_SUITE_TITLE, () => {
       describe(GRPC_TEST_SUITE_TITLE, () => {
@@ -753,6 +768,11 @@ describe('Tracing Tests', () => {
         SPAN_NAME_DOC_REF_GET,
         SPAN_NAME_BATCH_GET_DOCUMENTS
       );
+      expectSpanHasEvents(SPAN_NAME_BATCH_GET_DOCUMENTS, [
+        'Firestore.batchGetDocuments: Start',
+        'Firestore.batchGetDocuments: First response received',
+        'Firestore.batchGetDocuments: Completed',
+      ]);
     });
 
     it('document reference create()', async () => {
@@ -820,6 +840,11 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_AGGREGATION_QUERY_GET);
+      expectSpanHasEvents(SPAN_NAME_AGGREGATION_QUERY_GET, [
+        'Firestore.runAggregationQuery: Start',
+        'Firestore.runAggregationQuery: First response received',
+        'Firestore.runAggregationQuery: Completed',
+      ]);
     });
 
     it('collection reference add()', async () => {
@@ -852,6 +877,12 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_QUERY_GET);
+      expectSpanHasEvents(SPAN_NAME_QUERY_GET, [
+        'RunQuery',
+        'Firestore.runQuery: Start',
+        'Firestore.runQuery: First response received',
+        'Firestore.runQuery: Completed',
+      ]);
     });
 
     it('firestore getAll()', async () => {
@@ -862,6 +893,11 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_BATCH_GET_DOCUMENTS);
+      expectSpanHasEvents(SPAN_NAME_BATCH_GET_DOCUMENTS, [
+        'Firestore.batchGetDocuments: Start',
+        'Firestore.batchGetDocuments: First response received',
+        'Firestore.batchGetDocuments: Completed',
+      ]);
     });
 
     it('transaction', async () => {
@@ -920,6 +956,7 @@ describe('Tracing Tests', () => {
       await runFirestoreOperationInRootSpan(async () => {
         const query = firestore.collectionGroup('foo');
         let numPartitions = 0;
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
         for await (const partition of query.getPartitions(3)) {
           numPartitions++;
         }