refactor: Add settings attributes to trace spans. (#2181)

ehsannas · web-flow · commit 1775ce6e9a32 · 2024-09-23T11:40:02.000-07:00
diff --git a/dev/src/index.ts b/dev/src/index.ts
@@ -179,7 +179,7 @@ export const DEFAULT_MAX_TRANSACTION_ATTEMPTS = 5;
 /*!
  * The default number of idle GRPC channel to keep.
  */
-const DEFAULT_MAX_IDLE_CHANNELS = 1;
+export const DEFAULT_MAX_IDLE_CHANNELS = 1;
 
 /*!
  * The maximum number of concurrent requests supported by a single GRPC channel,
@@ -589,7 +589,6 @@ export class Firestore implements firestore.Firestore {
     }
 
     this.validateAndApplySettings({...settings, ...libraryHeader});
-
     this._traceUtil = this.newTraceUtilInstance(this._settings);
 
     const retryConfig = serviceConfig.retry_params.default;
@@ -696,6 +695,7 @@ export class Firestore implements firestore.Firestore {
 
     const mergedSettings = {...this._settings, ...settings};
     this.validateAndApplySettings(mergedSettings);
+    this._traceUtil = this.newTraceUtilInstance(this._settings);
     this._settingsFrozen = true;
   }
 
@@ -789,7 +789,6 @@ export class Firestore implements firestore.Firestore {
       return temp;
     };
     this._serializer = new Serializer(this);
-    this._traceUtil = this.newTraceUtilInstance(this._settings);
   }
 
   private newTraceUtilInstance(settings: firestore.Settings): TraceUtil {
@@ -809,10 +808,6 @@ export class Firestore implements firestore.Firestore {
     }
 
     if (createEnabledInstance) {
-      // Re-use the existing EnabledTraceUtil if one has been created.
-      if (this._traceUtil && this._traceUtil instanceof EnabledTraceUtil) {
-        return this._traceUtil;
-      }
       return new EnabledTraceUtil(settings);
     } else {
       return new DisabledTraceUtil();
@@ -1550,6 +1545,10 @@ export class Firestore implements firestore.Firestore {
           'Detected project ID: %s',
           this._projectId
         );
+
+        // If the project ID was undefined when the TraceUtil was set up, we
+        // need to record it.
+        this._traceUtil.recordProjectId(this.projectId);
       } catch (err) {
         logger(
           'Firestore.initializeIfNeeded',
diff --git a/dev/src/telemetry/disabled-trace-util.ts b/dev/src/telemetry/disabled-trace-util.ts
@@ -36,4 +36,7 @@ export class DisabledTraceUtil implements TraceUtil {
   currentSpan(): Span {
     return new Span();
   }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  recordProjectId(projectId: string): void {}
 }
diff --git a/dev/src/telemetry/enabled-trace-util.ts b/dev/src/telemetry/enabled-trace-util.ts
@@ -25,10 +25,17 @@ import {
 } from '@opentelemetry/api';
 
 import {Span} from './span';
-import {Attributes, TraceUtil} from './trace-util';
+import {ATTRIBUTE_SETTINGS_PREFIX, Attributes, TraceUtil} from './trace-util';
+
+import {interfaces} from '../v1/firestore_client_config.json';
+import {FirestoreClient} from '../v1';
+import {DEFAULT_DATABASE_ID} from '../path';
+import {DEFAULT_MAX_IDLE_CHANNELS} from '../index';
+const serviceConfig = interfaces['google.firestore.v1.Firestore'];
 
 export class EnabledTraceUtil implements TraceUtil {
   private tracer: Tracer;
+  private settingsAttributes: Attributes;
 
   constructor(settings: Settings) {
     let tracerProvider = settings.openTelemetryOptions?.tracerProvider;
@@ -42,6 +49,85 @@ export class EnabledTraceUtil implements TraceUtil {
     const libVersion = require('../../../package.json').version;
     const libName = require('../../../package.json').name;
     this.tracer = tracerProvider.getTracer(libName, libVersion);
+
+    this.settingsAttributes = {};
+    this.settingsAttributes['otel.scope.name'] = libName;
+    this.settingsAttributes['otel.scope.version'] = libVersion;
+
+    if (settings.projectId) {
+      this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.project_id`] =
+        settings.projectId;
+    }
+
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.database_id`] =
+      settings.databaseId || DEFAULT_DATABASE_ID;
+
+    const host =
+      settings.servicePath ?? settings.host ?? 'firestore.googleapis.com';
+    const port = settings.port ?? FirestoreClient.port;
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.host`] =
+      `${host}:${port}`;
+
+    if (settings.preferRest !== undefined) {
+      this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.prefer_REST`] =
+        settings.preferRest;
+    }
+
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_idle_channels`] =
+      settings.maxIdleChannels ?? DEFAULT_MAX_IDLE_CHANNELS;
+
+    const defaultRetrySettings = serviceConfig.retry_params.default;
+    const customRetrySettings =
+      settings.clientConfig?.interfaces?.['google.firestore.v1.Firestore']?.[
+        'retry_params'
+      ]?.['default'];
+    this.settingsAttributes[
+      `${ATTRIBUTE_SETTINGS_PREFIX}.initial_retry_delay`
+    ] = this.millisToSecondString(
+      customRetrySettings?.initial_retry_delay_millis ??
+        defaultRetrySettings.initial_retry_delay_millis
+    );
+    this.settingsAttributes[
+      `${ATTRIBUTE_SETTINGS_PREFIX}.initial_rpc_timeout`
+    ] = this.millisToSecondString(
+      customRetrySettings?.initial_rpc_timeout_millis ??
+        defaultRetrySettings.initial_rpc_timeout_millis
+    );
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.total_timeout`] =
+      this.millisToSecondString(
+        customRetrySettings?.total_timeout_millis ??
+          defaultRetrySettings.total_timeout_millis
+      );
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_retry_delay`] =
+      this.millisToSecondString(
+        customRetrySettings?.max_retry_delay_millis ??
+          defaultRetrySettings.max_retry_delay_millis
+      );
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_rpc_timeout`] =
+      this.millisToSecondString(
+        customRetrySettings?.max_rpc_timeout_millis ??
+          defaultRetrySettings.max_rpc_timeout_millis
+      );
+    this.settingsAttributes[
+      `${ATTRIBUTE_SETTINGS_PREFIX}.retry_delay_multiplier`
+    ] =
+      customRetrySettings?.retry_delay_multiplier.toString() ??
+      defaultRetrySettings.retry_delay_multiplier.toString();
+    this.settingsAttributes[
+      `${ATTRIBUTE_SETTINGS_PREFIX}.rpc_timeout_multiplier`
+    ] =
+      customRetrySettings?.rpc_timeout_multiplier.toString() ??
+      defaultRetrySettings.rpc_timeout_multiplier.toString();
+  }
+
+  recordProjectId(projectId: string): void {
+    this.settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.project_id`] =
+      projectId;
+    this.currentSpan().setAttributes(this.settingsAttributes);
+  }
+
+  private millisToSecondString(millis: number): string {
+    return `${millis / 1000}s`;
   }
 
   private endSpan(otelSpan: OpenTelemetrySpan, error: Error): void {
@@ -64,6 +150,8 @@ export class EnabledTraceUtil implements TraceUtil {
         attributes: attributes,
       },
       (otelSpan: OpenTelemetrySpan) => {
+        this.addCommonAttributes(otelSpan);
+
         // Note that if `fn` returns a `Promise`, we want the otelSpan to end
         // after the `Promise` has resolved, NOT after the `fn` has returned.
         // Therefore, we should not use a `finally` clause to end the otelSpan.
@@ -94,10 +182,16 @@ export class EnabledTraceUtil implements TraceUtil {
   }
 
   startSpan(name: string): Span {
-    return new Span(this.tracer.startSpan(name, undefined, context.active()));
+    const otelSpan = this.tracer.startSpan(name, undefined, context.active());
+    this.addCommonAttributes(otelSpan);
+    return new Span(otelSpan);
   }
 
   currentSpan(): Span {
     return new Span(trace.getActiveSpan());
   }
+
+  addCommonAttributes(otelSpan: OpenTelemetrySpan): void {
+    otelSpan.setAttributes(this.settingsAttributes);
+  }
 }
diff --git a/dev/src/telemetry/trace-util.ts b/dev/src/telemetry/trace-util.ts
@@ -57,6 +57,8 @@ export const SPAN_NAME_TRANSACTION_COMMIT = 'Transaction.Commit';
 export const SPAN_NAME_BATCH_COMMIT = 'Batch.Commit';
 export const SPAN_NAME_PARTITION_QUERY = 'PartitionQuery';
 export const SPAN_NAME_BULK_WRITER_COMMIT = 'BulkWriter.Commit';
+export const ATTRIBUTE_SERVICE_PREFIX = 'gcp.firestore';
+export const ATTRIBUTE_SETTINGS_PREFIX = `${ATTRIBUTE_SERVICE_PREFIX}.settings`;
 export const ATTRIBUTE_KEY_DOC_COUNT = 'doc_count';
 export const ATTRIBUTE_KEY_IS_TRANSACTIONAL = 'transactional';
 export const ATTRIBUTE_KEY_NUM_RESPONSES = 'response_count';
@@ -74,4 +76,6 @@ export interface TraceUtil {
   startSpan(name: string): Span;
 
   currentSpan(): Span;
+
+  recordProjectId(projectId: string): void;
 }
diff --git a/dev/system-test/tracing.ts b/dev/system-test/tracing.ts
@@ -44,6 +44,7 @@ import {setLogFunction, Firestore} from '../src';
 import {verifyInstance} from '../test/util/helpers';
 import {
   ATTRIBUTE_KEY_DOC_COUNT,
+  ATTRIBUTE_SETTINGS_PREFIX,
   SPAN_NAME_AGGREGATION_QUERY_GET,
   SPAN_NAME_BATCH_COMMIT,
   SPAN_NAME_BATCH_GET_DOCUMENTS,
@@ -66,7 +67,6 @@ import {
   SPAN_NAME_TRANSACTION_RUN,
 } from '../src/telemetry/trace-util';
 import {AsyncLocalStorageContextManager} from '@opentelemetry/context-async-hooks';
-import {deepStrictEqual} from 'assert';
 import {cloudtrace_v1, auth as gAuth} from '@googleapis/cloudtrace';
 import Schema$Trace = cloudtrace_v1.Schema$Trace;
 import Schema$TraceSpan = cloudtrace_v1.Schema$TraceSpan;
@@ -306,6 +306,43 @@ describe('Tracing Tests', () => {
     firestore = new Firestore(settings);
   }
 
+  function getSettingsAttributes(): Attributes {
+    const settingsAttributes: Attributes = {};
+    settingsAttributes['otel.scope.name'] = require('../../package.json').name;
+    settingsAttributes['otel.scope.version'] =
+      require('../../package.json').version;
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.database_id`] =
+      firestore.databaseId;
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.host`] =
+      'firestore.googleapis.com:443';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.prefer_REST`] =
+      testConfig.preferRest;
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_idle_channels`] = 1;
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.initial_retry_delay`] =
+      '0.1s';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.initial_rpc_timeout`] =
+      '60s';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.total_timeout`] = '600s';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_retry_delay`] = '60s';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.max_rpc_timeout`] = '60s';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.retry_delay_multiplier`] =
+      '1.3';
+    settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.rpc_timeout_multiplier`] =
+      '1';
+
+    // Project ID is not set on the Firestore object until _after_ the first
+    // operation is done. Therefore, the spans that are created _before_ the
+    // first operation do not contain a project ID.
+    try {
+      const projectId = firestore.projectId;
+      settingsAttributes[`${ATTRIBUTE_SETTINGS_PREFIX}.project_id`] = projectId;
+    } catch (e) {
+      // Project ID has not been set yet.
+    }
+
+    return settingsAttributes;
+  }
+
   // Take a function and runs it inside a new root span. This makes it possible to
   // encapsulate all the SDK-generated spans inside a test root span. It also makes
   // it easy to query a trace storage backend for a known trace ID and span Id.
@@ -588,7 +625,26 @@ describe('Tracing Tests', () => {
         parentSpan.traceId,
         `'${childSpan.name}' and '${parentSpan.name}' spans do not belong to the same trace`
       );
-      // TODO(tracing): expect that each span has the needed attributes.
+
+      // The Cloud Trace API does not return span attributes and events.
+      if (!testConfig.e2e) {
+        const settingsAttributes = getSettingsAttributes();
+        for (const attributesKey in settingsAttributes) {
+          if (
+            attributesKey.endsWith('.project_id') &&
+            i + 1 !== matchingSpanHierarchy.length
+          ) {
+            // Project ID is not set on the Firestore object until _after_ the first
+            // operation is done. Therefore, the spans that are created _before_ the
+            // first operation do not contain a project ID. So, we'll just compare
+            // this attribute on the leaf spans.
+          } else {
+            expect(childSpan.attributes[attributesKey]).to.be.equal(
+              settingsAttributes[attributesKey]
+            );
+          }
+        }
+      }
     }
   }
 
@@ -597,7 +653,7 @@ describe('Tracing Tests', () => {
     spanName: string,
     attributes: Attributes
   ): void {
-    // TODO(tracing): The current Cloud Trace API does not return span attributes and events.
+    // The Cloud Trace API does not return span attributes and events.
     if (testConfig.e2e) {
       return;
     }
@@ -606,8 +662,14 @@ describe('Tracing Tests', () => {
     const span = getSpanByName(spanName);
     expect(span).to.not.be.null;
 
-    // Assert that the attributes are the same.
-    deepStrictEqual(span!.attributes, attributes);
+    // 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.
+    for (const attributesKey in attributes) {
+      expect(span!.attributes[attributesKey]).to.be.equal(
+        attributes[attributesKey]
+      );
+    }
   }
 
   describe(IN_MEMORY_TEST_SUITE_TITLE, () => {

Original file line number	Diff line number	Diff line change
`@@ -36,4 +36,7 @@ export class DisabledTraceUtil implements TraceUtil {`
`36`	`36`	`currentSpan(): Span {`
`37`	`37`	`return new Span();`
`38`	`38`	`}`
	`39`	`+`
	`40`	`+ // eslint-disable-next-line @typescript-eslint/no-unused-vars`
	`41`	`+ recordProjectId(projectId: string): void {}`
`39`	`42`	`}`