Skip to content

Commit

Permalink
refactor: Add settings attributes to trace spans.
Browse files Browse the repository at this point in the history
  • Loading branch information
@ehsannas
ehsannas committed Sep 4, 2024
1 parent 4fcefde commit 000432f
Show file tree
Hide file tree
Showing 5 changed files with 176 additions and 14 deletions.
13 changes: 6 additions & 7 deletions dev/src/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,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,
Expand Down Expand Up @@ -588,7 +588,6 @@ export class Firestore implements firestore.Firestore {
}

this.validateAndApplySettings({...settings, ...libraryHeader});

this._traceUtil = this.newTraceUtilInstance(this._settings);

const retryConfig = serviceConfig.retry_params.default;
Expand Down Expand Up @@ -695,6 +694,7 @@ export class Firestore implements firestore.Firestore {

const mergedSettings = {...this._settings, ...settings};
this.validateAndApplySettings(mergedSettings);
this._traceUtil = this.newTraceUtilInstance(this._settings);
this._settingsFrozen = true;
}

Expand Down Expand Up @@ -788,7 +788,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 {
Expand All @@ -808,10 +807,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();
Expand Down Expand Up @@ -1549,6 +1544,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',
Expand Down
3 changes: 3 additions & 0 deletions dev/src/telemetry/disabled-trace-util.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}
}
98 changes: 96 additions & 2 deletions dev/src/telemetry/enabled-trace-util.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand All @@ -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.
Expand Down Expand Up @@ -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);
}
}
4 changes: 4 additions & 0 deletions dev/src/telemetry/trace-util.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -74,4 +76,6 @@ export interface TraceUtil {
startSpan(name: string): Span;

currentSpan(): Span;

recordProjectId(projectId: string): void;
}
72 changes: 67 additions & 5 deletions dev/system-test/tracing.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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]
);
}
}
}
}
}

Expand All @@ -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;
}
Expand All @@ -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, () => {
Expand Down

0 comments on commit 000432f

Please sign in to comment.