|
1 | 1 | import { Hub } from '@sentry/core'; |
2 | 2 | import { EventProcessor, Integration, SpanContext } from '@sentry/types'; |
3 | | -import { fill, isThenable, loadModule, logger } from '@sentry/utils'; |
| 3 | +import { fill, isInstanceOf, isThenable, loadModule, logger } from '@sentry/utils'; |
| 4 | +import { EventEmitter } from 'events'; |
4 | 5 |
|
5 | 6 | import { shouldDisableAutoInstrumentation } from './utils/node-utils'; |
6 | 7 |
|
@@ -160,20 +161,38 @@ export class Mongo implements Integration { |
160 | 161 | // its (non-callback) arguments can also be functions.) |
161 | 162 | if (typeof lastArg !== 'function' || (operation === 'mapReduce' && args.length === 2)) { |
162 | 163 | const span = parentSpan?.startChild(getSpanContext(this, operation, args)); |
163 | | - const maybePromise = orig.call(this, ...args) as Promise<unknown>; |
| 164 | + const maybePromiseOrCursor = orig.call(this, ...args); |
164 | 165 |
|
165 | | - if (isThenable(maybePromise)) { |
166 | | - return maybePromise.then((res: unknown) => { |
| 166 | + if (isThenable(maybePromiseOrCursor)) { |
| 167 | + return maybePromiseOrCursor.then((res: unknown) => { |
167 | 168 | span?.finish(); |
168 | 169 | return res; |
169 | 170 | }); |
| 171 | + } |
| 172 | + // If the operation returns a cursor (which is an EventEmitter), |
| 173 | + // we need to attach a listener to it to finish the span when the cursor is closed. |
| 174 | + else if (isInstanceOf(maybePromiseOrCursor, EventEmitter)) { |
| 175 | + const cursor = maybePromiseOrCursor as EventEmitter; |
| 176 | + |
| 177 | + try { |
| 178 | + cursor.once('close', () => { |
| 179 | + span?.finish(); |
| 180 | + }); |
| 181 | + } catch (e) { |
| 182 | + // If the cursor is already closed, `once` will throw an error. In that case, we can |
| 183 | + // finish the span immediately. |
| 184 | + span?.finish(); |
| 185 | + } |
| 186 | + |
| 187 | + return cursor; |
170 | 188 | } else { |
171 | 189 | span?.finish(); |
172 | | - return maybePromise; |
| 190 | + return maybePromiseOrCursor; |
173 | 191 | } |
174 | 192 | } |
175 | 193 |
|
176 | 194 | const span = parentSpan?.startChild(getSpanContext(this, operation, args.slice(0, -1))); |
| 195 | + |
177 | 196 | return orig.call(this, ...args.slice(0, -1), function (err: Error, result: unknown) { |
178 | 197 | span?.finish(); |
179 | 198 | lastArg(err, result); |
|
0 commit comments