diff --git a/observability-test/spanner.ts b/observability-test/spanner.ts index bf3d93538..fc053d566 100644 --- a/observability-test/spanner.ts +++ b/observability-test/spanner.ts @@ -37,6 +37,9 @@ const { const {ObservabilityOptions} = require('../src/instrument'); +const selectSql = 'SELECT 1'; +const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2'; + /** A simple result set for SELECT 1. */ function createSelect1ResultSet(): protobuf.ResultSet { const fields = [ @@ -141,8 +144,6 @@ describe('EndToEnd', () => { server = setupResult.server; spannerMock = setupResult.spannerMock; - const selectSql = 'SELECT 1'; - const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2'; spannerMock.putStatementResult( selectSql, mock.StatementResult.resultSet(createSelect1ResultSet()) @@ -223,7 +224,13 @@ describe('EndToEnd', () => { `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` ); - const expectedEventNames = []; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Begin Transaction', + 'Transaction Creation Done', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -260,7 +267,12 @@ describe('EndToEnd', () => { `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` ); - const expectedEventNames = ['Using Session']; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Using Session', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -297,7 +309,12 @@ describe('EndToEnd', () => { `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` ); - const expectedEventNames = ['Using Session']; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Using Session', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -359,7 +376,12 @@ describe('EndToEnd', () => { 'Expected that RunStream has a defined spanId' ); - const expectedEventNames = ['Using Session']; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Using Session', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -393,7 +415,11 @@ describe('EndToEnd', () => { `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` ); - const expectedEventNames = []; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -431,7 +457,14 @@ describe('EndToEnd', () => { `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` ); - const expectedEventNames = ['Using Session']; + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Using Session', + 'Starting Commit', + 'Commit Done', + ]; assert.deepStrictEqual( actualEventNames, expectedEventNames, @@ -583,3 +616,118 @@ describe('ObservabilityOptions injection and propagation', async () => { }); }); }); + +describe('Bug fixes', () => { + it('async/await correctly parents trace spans', async () => { + const traceExporter = new InMemorySpanExporter(); + const provider = new NodeTracerProvider({ + sampler: new AlwaysOnSampler(), + exporter: traceExporter, + }); + provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter)); + + const observabilityOptions: typeof ObservabilityOptions = { + tracerProvider: provider, + enableExtendedTracing: true, + }; + const setupResult = await setup(observabilityOptions); + const spanner = setupResult.spanner; + const server = setupResult.server; + const spannerMock = setupResult.spannerMock; + + after(async () => { + provider.shutdown(); + spannerMock.resetRequests(); + spanner.close(); + server.tryShutdown(() => {}); + }); + + async function main() { + const instance = spanner.instance('testing'); + instance._observabilityOptions = observabilityOptions; + const database = instance.database('db-1'); + + const query = { + sql: selectSql, + }; + + const [rows] = await database.run(query); + + rows.forEach(row => { + const json = row.toJSON(); + }); + + provider.forceFlush(); + } + + await main(); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + + const actualSpanNames: string[] = []; + const actualEventNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + span.events.forEach(event => { + actualEventNames.push(event.name); + }); + }); + + const expectedSpanNames = [ + 'CloudSpanner.Database.runStream', + 'CloudSpanner.Database.run', + ]; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + // We need to ensure a strict relationship between the spans. + // runSpan -------------------| + // |-runStream ----------| + const runStreamSpan = spans[0]; + const runSpan = spans[1]; + assert.ok( + runSpan.spanContext().traceId, + 'Expected that unSpan has a defined traceId' + ); + assert.ok( + runStreamSpan.spanContext().traceId, + 'Expected that runStreamSpan has a defined traceId' + ); + + assert.deepStrictEqual( + runStreamSpan.parentSpanId, + runSpan.spanContext().spanId, + `Expected that runSpan(spanId=${runSpan.spanContext().spanId}) is the parent to runStreamSpan(parentSpanId=${runStreamSpan.parentSpanId})` + ); + + assert.deepStrictEqual( + runSpan.spanContext().traceId, + runStreamSpan.spanContext().traceId, + 'Expected that both spans share a traceId' + ); + assert.ok( + runStreamSpan.spanContext().spanId, + 'Expected that runStreamSpan has a defined spanId' + ); + assert.ok( + runSpan.spanContext().spanId, + 'Expected that runSpan has a defined spanId' + ); + + const expectedEventNames = [ + 'Acquiring session', + 'Waiting for a session to become available', + 'Acquired session', + 'Using Session', + ]; + assert.deepStrictEqual( + actualEventNames, + expectedEventNames, + `Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}` + ); + }); +}); diff --git a/src/instrument.ts b/src/instrument.ts index 6cad7bc4a..b81b0a9c5 100644 --- a/src/instrument.ts +++ b/src/instrument.ts @@ -26,7 +26,9 @@ import { import { Span, SpanStatusCode, + context, trace, + ROOT_CONTEXT, INVALID_SPAN_CONTEXT, SpanAttributes, TimeInput, @@ -93,6 +95,30 @@ interface traceConfig { const SPAN_NAMESPACE_PREFIX = 'CloudSpanner'; // TODO: discuss & standardize this prefix. export {SPAN_NAMESPACE_PREFIX, traceConfig}; +const { + AsyncHooksContextManager, +} = require('@opentelemetry/context-async-hooks'); + +/* + * This function ensures that async/await functions correctly by + * checking if context.active() returns an invalid/unset context + * and if so, sets a global AsyncHooksContextManager. + */ +function ensureInitialContextManagerSet() { + if (context.active() === ROOT_CONTEXT) { + // If no active context was set previously, trace context propagation cannot + // work correctly with async/await for OpenTelemetry and they acknowledge + // this fact per https://opentelemetry.io/docs/languages/js/context/#active-context + // but we shouldn't make our customers have to invasively edit their code + // nor should they be burdened about these facts, their code should JUST work. + // Please see https://github.com/googleapis/nodejs-spanner/issues/2146 + context.disable(); // Firstly disable any prior contextManager. + const contextManager = new AsyncHooksContextManager(); + contextManager.enable(); + context.setGlobalContextManager(contextManager); + } +} + /** * startTrace begins an active span in the current active context * and passes it back to the set callback function. Each span will @@ -111,6 +137,8 @@ export function startTrace( config = {} as traceConfig; } + ensureInitialContextManagerSet(); + return getTracer(config.opts?.tracerProvider).startActiveSpan( SPAN_NAMESPACE_PREFIX + '.' + spanNameSuffix, {kind: SpanKind.CLIENT},