Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: (observability) make async/await correctly work by setting initial AsyncHooksManager #2147

Closed
wants to merge 2 commits into from
+187 −8
Closed

fix: (observability) make async/await correctly work by setting initial AsyncHooksManager #2147

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
44bc1b7
fix: (observability) make async/await correctly work by setting initi…
odeke-em Oct 4, 2024
9d04d6d
Tweak documentation and test name
odeke-em Oct 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 157 additions & 8 deletions observability-test/spanner.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -583,3 +616,119 @@ describe('ObservabilityOptions injection and propagation', async () => {
});
});
});

describe('Regression tests for fixed bugs', () => {
it('async/await correctly parents trace spans', async () => {
// See https://github.com/googleapis/nodejs-spanner/issues/2146.
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}`
);
});
});
30 changes: 30 additions & 0 deletions src/instrument.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,9 @@ import {
import {
Span,
SpanStatusCode,
context,
trace,
ROOT_CONTEXT,
INVALID_SPAN_CONTEXT,
SpanAttributes,
TimeInput,
Expand Down Expand Up @@ -93,6 +95,32 @@ 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 works correctly by
* checking if context.active() returns an invalid/unset context
* and if so, sets a global AsyncHooksContextManager otherwise
* spans resulting from async/await invocations won't be correctly
* associated in their respective hierarchies.
*/
function ensureInitialContextManagerSet() {
if (context.active() === ROOT_CONTEXT) {
// If no active context was set previously, trace context propagation cannot
// function 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
Expand All @@ -111,6 +139,8 @@ export function startTrace<T>(
config = {} as traceConfig;
}

ensureInitialContextManagerSet();

return getTracer(config.opts?.tracerProvider).startActiveSpan(
SPAN_NAMESPACE_PREFIX + '.' + spanNameSuffix,
{kind: SpanKind.CLIENT},
Expand Down