chore: fixing context propagation on mongo callback (open-telemetry#403)

Co-authored-by: Valentin Marchaud <[email protected]>

examples/mongodb/package.json

@@ -33,8 +33,8 @@
  "@opentelemetry/exporter-zipkin": "^0.18.0",
  "@opentelemetry/instrumentation": "^0.18.0",
  "@opentelemetry/instrumentation-http": "^0.18.0",
+ "@opentelemetry/instrumentation-mongodb": "^0.15.0",
  "@opentelemetry/node": "^0.18.0",
- "@opentelemetry/plugin-mongodb": "^0.14.0",
  "@opentelemetry/tracing": "^0.18.0",
  "mongodb": "^3.5.7"
  },

examples/mongodb/server.js

@@ -60,13 +60,17 @@ startServer(8080);
 
 function handleInsertQuery(response) {
  const obj = { name: 'John', age: '20' };
- db.collection('users').insertOne(obj, (err) => {
+ const collection = db.collection('users');
+ collection.insertOne(obj, (err) => {
  if (err) {
  console.log('Error code:', err.code);
  response.end(err.message);
  } else {
  console.log('1 document inserted');
- response.end();
+ // find document to test context propagation using callback
+ collection.findOne({}, function () {
+ response.end();
+ });
  }
  });
 }

examples/mongodb/tracer.js

@@ -7,6 +7,7 @@ const { JaegerExporter } = require('@opentelemetry/exporter-jaeger');
 const { ZipkinExporter } = require('@opentelemetry/exporter-zipkin');
 const { registerInstrumentations } = require('@opentelemetry/instrumentation');
 const { HttpInstrumentation } = require('@opentelemetry/instrumentation-http');
+const { MongoDBInstrumentation } = require('@opentelemetry/instrumentation-mongodb');
 
 module.exports = (serviceName) => {
  const provider = new NodeTracerProvider();
@@ -24,13 +25,12 @@ module.exports = (serviceName) => {
  registerInstrumentations({
  instrumentations: [
  new HttpInstrumentation(),
+ new MongoDBInstrumentation({
+ enhancedDatabaseReporting: true,
+ }),
  {
  plugins: {
- mongodb: {
- enabled: true,
- path: '@opentelemetry/plugin-mongodb',
- enhancedDatabaseReporting: true,
- },
+ mongodb: { enabled: false },
  },
  },
  ],

plugins/node/opentelemetry-instrumentation-mongodb/package.json

@@ -6,6 +6,7 @@
  "types": "build/src/index.d.ts",
  "repository": "open-telemetry/opentelemetry-js-contrib",
  "scripts": {
+ "docker:start": "docker run -e MONGODB_DB=opentelemetry-tests -e MONGODB_PORT=27017 -e MONGODB_HOST=localhost -p 27017:27017 --rm mongo",
  "test": "nyc ts-mocha --parallel -p tsconfig.json 'test/**/*.test.ts'",
  "codecov": "nyc report --reporter=json && codecov -f coverage/*.json -p ../../",
  "tdd": "npm run test -- --watch-extensions ts --watch",
@@ -15,7 +16,8 @@
  "precompile": "tsc --version",
  "version:update": "node ../../../scripts/version-update.js",
  "compile": "npm run version:update && tsc -p .",
- "prepare": "npm run compile"
+ "prepare": "npm run compile",
+ "watch": "tsc -w"
  },
  "keywords": [
  "opentelemetry",

plugins/node/opentelemetry-instrumentation-mongodb/src/mongodb.ts

@@ -397,6 +397,9 @@ export class MongoDBInstrumentation extends InstrumentationBase<
  * @param resultHandler A callback function.
  */
  private _patchEnd(span: Span, resultHandler: Function): Function {
+ // mongodb is using "tick" when calling a callback, this way the context
+ // in final callback (resultHandler) is lost
+ const activeContext = context.active();
  return function patchedEnd(this: {}, ...args: unknown[]) {
  const error = args[0];
  if (error instanceof Error) {
@@ -406,7 +409,10 @@ export class MongoDBInstrumentation extends InstrumentationBase<
  });
  }
  span.end();
- return resultHandler.apply(this, args);
+
+ return context.with(activeContext, () => {
+ return resultHandler.apply(this, args);
+ });
  };
  }
 }

plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts

@@ -14,8 +14,7 @@
  * limitations under the License.
  */
 
-// for testing locally use this command to run docker
-// docker run -e MONGODB_DB=opentelemetry-tests -e MONGODB_PORT=27017 -e MONGODB_HOST=localhost -p 27017:27017 --name otmongo mongo
+// for testing locally "npm run docker:start"
 
 import { context, setSpan, SpanKind } from '@opentelemetry/api';
 import { BasicTracerProvider } from '@opentelemetry/tracing';
@@ -99,6 +98,7 @@ describe('MongoDBInstrumentation', () => {
 
  afterEach(done => {
  collection.deleteMany({}, done);
+ memoryExporter.reset();
  });
 
  after(() => {
@@ -109,6 +109,9 @@ describe('MongoDBInstrumentation', () => {
 
  /** Should intercept query */
  describe('Instrumenting query operations', () => {
+ beforeEach(() => {
+ memoryExporter.reset();
+ });
  it('should create a child span for insert', done => {
  const insertData = [{ a: 1 }, { a: 2 }, { a: 3 }];
  const span = provider.getTracer('default').startSpan('insertRootSpan');
@@ -161,6 +164,10 @@ describe('MongoDBInstrumentation', () => {
 
  /** Should intercept cursor */
  describe('Instrumenting cursor operations', () => {
+ beforeEach(() => {
+ memoryExporter.reset();
+ });
+
  it('should create a child span for find', done => {
  const span = provider.getTracer('default').startSpan('findRootSpan');
  context.with(setSpan(context.active(), span), () => {
@@ -212,6 +219,10 @@ describe('MongoDBInstrumentation', () => {
 
  /** Should intercept command */
  describe('Instrumenting command operations', () => {
+ beforeEach(() => {
+ memoryExporter.reset();
+ });
+
  it('should create a child span for create index', done => {
  const span = provider.getTracer('default').startSpan('indexRootSpan');
  context.with(setSpan(context.active(), span), () => {
@@ -229,8 +240,47 @@ describe('MongoDBInstrumentation', () => {
  });
  });
 
+ describe('Mixed operations with callback', () => {
+ beforeEach(() => {
+ memoryExporter.reset();
+ });
+
+ it('should create a span for find after callback insert', done => {
+ const insertData = [{ a: 1 }, { a: 2 }, { a: 3 }];
+ const span = provider.getTracer('default').startSpan('insertRootSpan');
+ context.with(setSpan(context.active(), span), () => {
+ collection.insertMany(insertData, (err, result) => {
+ span.end();
+ assert.ifError(err);
+ const spans = memoryExporter.getFinishedSpans();
+ const mainSpan = spans[spans.length - 1];
+ assertSpans(spans, 'mongodb.insert', SpanKind.CLIENT);
+ memoryExporter.reset();
+
+ collection.find({ a: 1 }).toArray((err, result) => {
+ const spans2 = memoryExporter.getFinishedSpans();
+ spans2.push(mainSpan);
+
+ assert.ifError(err);
+ assertSpans(spans2, 'mongodb.find', SpanKind.CLIENT);
+ assert.strictEqual(
+ mainSpan.spanContext.spanId,
+ spans2[0].parentSpanId
+ );
+ memoryExporter.reset();
+ done();
+ });
+ });
+ });
+ });
+ });
+
  /** Should intercept command */
  describe('Removing Instrumentation', () => {
+ beforeEach(() => {
+ memoryExporter.reset();
+ });
+
  it('should unpatch plugin', () => {
  assert.doesNotThrow(() => {
  instrumentation.disable();