fix(tracing): Instrument cursors returned from MongoDB operations. (#6437)

Author: onurtemizkan

packages/node-integration-tests/suites/tracing/auto-instrument/mongodb/scenario.ts

@@ -36,6 +36,8 @@ async function run(): Promise<void> {
     await collection.findOne({ title: 'Back to the Future' });
     await collection.updateOne({ title: 'Back to the Future' }, { $set: { title: 'South Park' } });
     await collection.findOne({ title: 'South Park' });
+
+    await collection.find({ title: 'South Park' }).toArray();
   } finally {
     if (transaction) transaction.finish();
     await client.close();

packages/node-integration-tests/suites/tracing/auto-instrument/mongodb/test.ts

@@ -48,16 +48,6 @@ conditionalTest({ min: 12 })('MongoDB Test', () => {
           description: 'findOne',
           op: 'db',
         },
-        {
-          data: {
-            collectionName: 'movies',
-            dbName: 'admin',
-            namespace: 'admin.movies',
-            query: '{"title":"Back to the Future"}',
-          },
-          description: 'find',
-          op: 'db',
-        },
         {
           data: {
             collectionName: 'movies',

packages/tracing/src/integrations/node/mongo.ts

@@ -90,6 +90,14 @@ interface MongoOptions {
   useMongoose?: boolean;
 }
 
+interface MongoCursor {
+  once(event: 'close', listener: () => void): void;
+}
+
+function isCursor(maybeCursor: MongoCursor): maybeCursor is MongoCursor {
+  return maybeCursor && typeof maybeCursor === 'object' && maybeCursor.once && typeof maybeCursor.once === 'function';
+}
+
 /** Tracing integration for mongo package */
 export class Mongo implements Integration {
   /**
@@ -160,20 +168,38 @@ export class Mongo implements Integration {
         // its (non-callback) arguments can also be functions.)
         if (typeof lastArg !== 'function' || (operation === 'mapReduce' && args.length === 2)) {
           const span = parentSpan?.startChild(getSpanContext(this, operation, args));
-          const maybePromise = orig.call(this, ...args) as Promise<unknown>;
+          const maybePromiseOrCursor = orig.call(this, ...args);
 
-          if (isThenable(maybePromise)) {
-            return maybePromise.then((res: unknown) => {
+          if (isThenable(maybePromiseOrCursor)) {
+            return maybePromiseOrCursor.then((res: unknown) => {
               span?.finish();
               return res;
             });
+          }
+          // If the operation returns a Cursor
+          // we need to attach a listener to it to finish the span when the cursor is closed.
+          else if (isCursor(maybePromiseOrCursor)) {
+            const cursor = maybePromiseOrCursor as MongoCursor;
+
+            try {
+              cursor.once('close', () => {
+                span?.finish();
+              });
+            } catch (e) {
+              // If the cursor is already closed, `once` will throw an error. In that case, we can
+              // finish the span immediately.
+              span?.finish();
+            }
+
+            return cursor;
           } else {
             span?.finish();
-            return maybePromise;
+            return maybePromiseOrCursor;
           }
         }
 
         const span = parentSpan?.startChild(getSpanContext(this, operation, args.slice(0, -1)));
+
         return orig.call(this, ...args.slice(0, -1), function (err: Error, result: unknown) {
           span?.finish();
           lastArg(err, result);