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

Co-authored-by: Valentin Marchaud <thisismac47@gmail.com>

Author: obecny

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();