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
Closed

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

Show file tree
Hide file tree
Changes from 1 commit
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
164 changes: 156 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,118 @@ describe('ObservabilityOptions injection and propagation', async () => {
});
});
});

describe('Bug fixes', () => {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use a better name for this. If you are testing "async/await" use name like Observability tests with async/await

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd highly encourage keeping such a section for known and fixed regressions with isolated test cases, otherwise eventually we'll just have a pile up of test cases that we don't have a single reference to.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the section in which we shall be inserting all regression tests and bug fixes.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tweaked it, please take a look.

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}`
);
});
});
28 changes: 28 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,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.
*/
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/*
* 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.
*/
/**
* Ensures that an appropriate Context Manager is set to support proper
* context propagation when using async/await, particularly in OpenTelemetry.
*
* OpenTelemetry's context system relies on `context.active()` to retrieve the
* current active context. However, if this returns the default `ROOT_CONTEXT`,
* it indicates that no context has been previously set. Without a valid context,
* async/await calls will fail to propagate context correctly, which can cause
* issues in distributed tracing and telemetry data collection.
*
* This function checks if the active context is invalid (i.e., `ROOT_CONTEXT`).
* If so, it resets the Context Manager by disabling any prior manager and then
* enabling a new instance of `AsyncHooksContextManager` to properly track async
* operations and context propagation.
*
* The goal is to ensure that users' code works seamlessly with OpenTelemetry,
* without requiring them to modify their existing code to manually manage context.
*
* Further details:
* - OpenTelemetry context documentation: https://opentelemetry.io/docs/languages/js/context/#active-context
* - Related issue: https://github.com/googleapis/nodejs-spanner/issues/2146
*/

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@odeke-em I generated this documentation using AI, please review once more. But we need a more structured documentation

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some of this AI generated documentation is not correct like

 * The goal is to ensure that users' code works seamlessly with OpenTelemetry, 
* without requiring them to modify their existing code to manually manage context.

it's nothing about manually managing context but setting a global context manager.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@surbhigarg92 the generated AI comments instead just is listing out what the function does but not really why and the purpose of the code which is to check if there was any already set contextManager and if not, set one.

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// 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
// No active context set, so context propagation won't work with async/await.
// Fix by disabling any prior context manager and setting a new AsyncHooksContextManager.

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 +137,8 @@ export function startTrace<T>(
config = {} as traceConfig;
}

ensureInitialContextManagerSet();

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