@@ -90,6 +90,14 @@ interface MongoOptions {
90
90
useMongoose ?: boolean ;
91
91
}
92
92
93
+ interface MongoCursor {
94
+ once ( event : 'close' , listener : ( ) => void ) : void ;
95
+ }
96
+
97
+ function isCursor ( maybeCursor : MongoCursor ) : maybeCursor is MongoCursor {
98
+ return maybeCursor && typeof maybeCursor === 'object' && maybeCursor . once && typeof maybeCursor . once === 'function' ;
99
+ }
100
+
93
101
/** Tracing integration for mongo package */
94
102
export class Mongo implements Integration {
95
103
/**
@@ -160,20 +168,38 @@ export class Mongo implements Integration {
160
168
// its (non-callback) arguments can also be functions.)
161
169
if ( typeof lastArg !== 'function' || ( operation === 'mapReduce' && args . length === 2 ) ) {
162
170
const span = parentSpan ?. startChild ( getSpanContext ( this , operation , args ) ) ;
163
- const maybePromise = orig . call ( this , ...args ) as Promise < unknown > ;
171
+ const maybePromiseOrCursor = orig . call ( this , ...args ) ;
164
172
165
- if ( isThenable ( maybePromise ) ) {
166
- return maybePromise . then ( ( res : unknown ) => {
173
+ if ( isThenable ( maybePromiseOrCursor ) ) {
174
+ return maybePromiseOrCursor . then ( ( res : unknown ) => {
167
175
span ?. finish ( ) ;
168
176
return res ;
169
177
} ) ;
178
+ }
179
+ // If the operation returns a Cursor
180
+ // we need to attach a listener to it to finish the span when the cursor is closed.
181
+ else if ( isCursor ( maybePromiseOrCursor ) ) {
182
+ const cursor = maybePromiseOrCursor as MongoCursor ;
183
+
184
+ try {
185
+ cursor . once ( 'close' , ( ) => {
186
+ span ?. finish ( ) ;
187
+ } ) ;
188
+ } catch ( e ) {
189
+ // If the cursor is already closed, `once` will throw an error. In that case, we can
190
+ // finish the span immediately.
191
+ span ?. finish ( ) ;
192
+ }
193
+
194
+ return cursor ;
170
195
} else {
171
196
span ?. finish ( ) ;
172
- return maybePromise ;
197
+ return maybePromiseOrCursor ;
173
198
}
174
199
}
175
200
176
201
const span = parentSpan ?. startChild ( getSpanContext ( this , operation , args . slice ( 0 , - 1 ) ) ) ;
202
+
177
203
return orig . call ( this , ...args . slice ( 0 , - 1 ) , function ( err : Error , result : unknown ) {
178
204
span ?. finish ( ) ;
179
205
lastArg ( err , result ) ;
0 commit comments