ref(nextjs): Use flush code from withSentry in all backend wrappers

packages/nextjs/src/config/wrappers/types.ts

@@ -1,4 +1,4 @@
-import type { Transaction } from '@sentry/types';
+import type { Transaction, WrappedFunction } from '@sentry/types';
 import type { NextApiRequest, NextApiResponse } from 'next';
 
 // The `NextApiHandler` and `WrappedNextApiHandler` types are the same as the official `NextApiHandler` type, except:
@@ -42,4 +42,4 @@ export type AugmentedNextApiResponse = NextApiResponse & {
 };
 
 export type ResponseEndMethod = AugmentedNextApiResponse['end'];
-export type WrappedResponseEndMethod = AugmentedNextApiResponse['end'];
+export type WrappedResponseEndMethod = AugmentedNextApiResponse['end'] & WrappedFunction;

packages/nextjs/src/config/wrappers/utils/responseEnd.ts

@@ -0,0 +1,69 @@
+import { flush } from '@sentry/node';
+import { Transaction } from '@sentry/types';
+import { fill, logger } from '@sentry/utils';
+import { ServerResponse } from 'http';
+
+import { ResponseEndMethod, WrappedResponseEndMethod } from '../types';
+
+/**
+ * Wrap `res.end()` so that it closes the transaction and flushes events before letting the request finish.
+ *
+ * Note: This wraps a sync method with an async method. While in general that's not a great idea in terms of keeping
+ * things in the right order, in this case it's safe, because the native `.end()` actually *is* (effectively) async, and
+ * its run actually *is* (literally) awaited, just manually so (which reflects the fact that the core of the
+ * request/response code in Node by far predates the introduction of `async`/`await`). When `.end()` is done, it emits
+ * the `prefinish` event, and only once that fires does request processing continue. See
+ * https://github.com/nodejs/node/commit/7c9b607048f13741173d397795bac37707405ba7.
+ *
+ * Also note: `res.end()` isn't called until *after* all response data and headers have been sent, so blocking inside of
+ * `end` doesn't delay data getting to the end user. See
+ * https://nodejs.org/api/http.html#responseenddata-encoding-callback.
+ *
+ * @param transaction The transaction tracing request handling
+ * @param res: The request's corresponding response
+ */
+export function autoEndTransactionOnResponseEnd(transaction: Transaction, res: ServerResponse): void {
+  const wrapEndMethod = (origEnd: ResponseEndMethod): WrappedResponseEndMethod => {
+    return async function sentryWrappedEnd(this: ServerResponse, ...args: unknown[]) {
+      await finishTransaction(transaction, this);
+      await flushQueue();
+
+      return origEnd.call(this, ...args);
+    };
+  };
+
+  // Prevent double-wrapping
+  if (!(res.end as WrappedResponseEndMethod).__sentry_original__) {
+    fill(res, 'end', wrapEndMethod);
+  }
+}
+
+/** Finish the given response's transaction and set HTTP status data */
+export async function finishTransaction(transaction: Transaction | undefined, res: ServerResponse): Promise<void> {
+  if (transaction) {
+    transaction.setHttpStatus(res.statusCode);
+
+    // If any open spans are set to finish when the response ends, it sets up a race condition between their `finish`
+    // calls and the transaction's `finish` call - and any spans which lose the race will get dropped from the
+    // transaction. To prevent this, push `transaction.finish` to the next event loop so that it's guaranteed to lose
+    // the race, and wait for it to be done before flushing events.
+    const transactionFinished: Promise<void> = new Promise(resolve => {
+      setImmediate(() => {
+        transaction.finish();
+        resolve();
+      });
+    });
+    await transactionFinished;
+  }
+}
+
+/** Flush the event queue to ensure that events get sent to Sentry before the response is finished and the lambda ends */
+export async function flushQueue(): Promise<void> {
+  try {
+    __DEBUG_BUILD__ && logger.log('Flushing events...');
+    await flush(2000);
+    __DEBUG_BUILD__ && logger.log('Done flushing events');
+  } catch (e) {
+    __DEBUG_BUILD__ && logger.log('Error while flushing events:\n', e);
+  }
+}

packages/nextjs/src/config/wrappers/withSentryAPI.ts

@@ -1,4 +1,4 @@
-import { captureException, flush, getCurrentHub, startTransaction } from '@sentry/node';
+import { captureException, getCurrentHub, startTransaction } from '@sentry/node';
 import { extractTraceparentData, hasTracingEnabled } from '@sentry/tracing';
 import {
   addExceptionMechanism,
@@ -11,14 +11,8 @@ import {
 import * as domain from 'domain';
 
 import { formatAsCode, nextLogger } from '../../utils/nextLogger';
-import type {
-  AugmentedNextApiRequest,
-  AugmentedNextApiResponse,
-  NextApiHandler,
-  ResponseEndMethod,
-  WrappedNextApiHandler,
-  WrappedResponseEndMethod,
-} from './types';
+import type { AugmentedNextApiRequest, AugmentedNextApiResponse, NextApiHandler, WrappedNextApiHandler } from './types';
+import { autoEndTransactionOnResponseEnd, finishTransaction, flushQueue } from './utils/responseEnd';
 
 /**
  * Wrap the given API route handler for tracing and error capturing. Thin wrapper around `withSentry`, which only
@@ -72,11 +66,6 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
     }
     req.__withSentry_applied__ = true;
 
-    // first order of business: monkeypatch `res.end()` so that it will wait for us to send events to sentry before it
-    // fires (if we don't do this, the lambda will close too early and events will be either delayed or lost)
-    // eslint-disable-next-line @typescript-eslint/unbound-method
-    res.end = wrapEndMethod(res.end);
-
     // use a domain in order to prevent scope bleed between requests
     const local = domain.create();
     local.add(req);
@@ -86,6 +75,7 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
     // return a value. In our case, all any of the codepaths return is a promise of `void`, but nextjs still counts on
     // getting that before it will finish the response.
     const boundHandler = local.bind(async () => {
+      let transaction;
       const currentScope = getCurrentHub().getScope();
 
       if (currentScope) {
@@ -121,7 +111,7 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
 
           const reqMethod = `${(req.method || 'GET').toUpperCase()} `;
 
-          const transaction = startTransaction(
+          transaction = startTransaction(
             {
               name: `${reqMethod}${reqPath}`,
               op: 'http.server',
@@ -137,9 +127,7 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
           );
           currentScope.setSpan(transaction);
 
-          // save a link to the transaction on the response, so that even if there's an error (landing us outside of
-          // the domain), we can still finish it (albeit possibly missing some scope data)
-          res.__sentryTransaction = transaction;
+          autoEndTransactionOnResponseEnd(transaction, res);
         }
       }
 
@@ -188,8 +176,10 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
         // Make sure we have a chance to finish the transaction and flush events to Sentry before the handler errors
         // out. (Apps which are deployed on Vercel run their API routes in lambdas, and those lambdas will shut down the
         // moment they detect an error, so it's important to get this done before rethrowing the error. Apps not
-        // deployed serverlessly will run into this cleanup function again in `res.end(), but it'll just no-op.)
-        await finishSentryProcessing(res);
+        // deployed serverlessly will run into this cleanup code again in `res.end(), but the transaction will already
+        // be finished and the queue will already be empty, so effectively it'll just no-op.)
+        await finishTransaction(transaction, res);
+        await flushQueue();
 
         // We rethrow here so that nextjs can do with the error whatever it would normally do. (Sometimes "whatever it
         // would normally do" is to allow the error to bubble up to the global handlers - another reason we need to mark
@@ -203,57 +193,3 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
     return boundHandler();
   };
 }
-
-/**
- * Wrap `res.end()` so that it closes the transaction and flushes events before letting the request finish.
- *
- * Note: This wraps a sync method with an async method. While in general that's not a great idea in terms of keeping
- * things in the right order, in this case it's safe, because the native `.end()` actually *is* async, and its run
- * actually *is* awaited, just manually so (which reflects the fact that the core of the request/response code in Node
- * by far predates the introduction of `async`/`await`). When `.end()` is done, it emits the `prefinish` event, and
- * only once that fires does request processing continue. See
- * https://github.com/nodejs/node/commit/7c9b607048f13741173d397795bac37707405ba7.
- *
- * @param origEnd The original `res.end()` method
- * @returns The wrapped version
- */
-function wrapEndMethod(origEnd: ResponseEndMethod): WrappedResponseEndMethod {
-  return async function newEnd(this: AugmentedNextApiResponse, ...args: unknown[]) {
-    await finishSentryProcessing(this);
-
-    return origEnd.call(this, ...args);
-  };
-}
-
-/**
- * Close the open transaction (if any) and flush events to Sentry.
- *
- * @param res The outgoing response for this request, on which the transaction is stored
- */
-async function finishSentryProcessing(res: AugmentedNextApiResponse): Promise<void> {
-  const { __sentryTransaction: transaction } = res;
-
-  if (transaction) {
-    transaction.setHttpStatus(res.statusCode);
-
-    // Push `transaction.finish` to the next event loop so open spans have a better chance of finishing before the
-    // transaction closes, and make sure to wait until that's done before flushing events
-    const transactionFinished: Promise<void> = new Promise(resolve => {
-      setImmediate(() => {
-        transaction.finish();
-        resolve();
-      });
-    });
-    await transactionFinished;
-  }
-
-  // Flush the event queue to ensure that events get sent to Sentry before the response is finished and the lambda
-  // ends. If there was an error, rethrow it so that the normal exception-handling mechanisms can apply.
-  try {
-    __DEBUG_BUILD__ && logger.log('Flushing events...');
-    await flush(2000);
-    __DEBUG_BUILD__ && logger.log('Done flushing events');
-  } catch (e) {
-    __DEBUG_BUILD__ && logger.log('Error while flushing events:\n', e);
-  }
-}

packages/nextjs/src/config/wrappers/wrapperUtils.ts

@@ -1,10 +1,12 @@
 import { captureException, getCurrentHub, startTransaction } from '@sentry/core';
 import { getActiveTransaction } from '@sentry/tracing';
 import { Transaction } from '@sentry/types';
-import { baggageHeaderToDynamicSamplingContext, extractTraceparentData, fill } from '@sentry/utils';
+import { baggageHeaderToDynamicSamplingContext, extractTraceparentData } from '@sentry/utils';
 import * as domain from 'domain';
 import { IncomingMessage, ServerResponse } from 'http';
 
+import { autoEndTransactionOnResponseEnd } from './utils/responseEnd';
+
 declare module 'http' {
   interface IncomingMessage {
     _sentryTransaction?: Transaction;
@@ -26,15 +28,6 @@ function setTransactionOnRequest(transaction: Transaction, req: IncomingMessage)
   req._sentryTransaction = transaction;
 }
 
-function autoEndTransactionOnResponseEnd(transaction: Transaction, res: ServerResponse): void {
-  fill(res, 'end', (originalEnd: ServerResponse['end']) => {
-    return function (this: unknown, ...endArguments: Parameters<ServerResponse['end']>) {
-      transaction.finish();
-      return originalEnd.call(this, ...endArguments);
-    };
-  });
-}
-
 /**
  * Wraps a function that potentially throws. If it does, the error is passed to `captureException` and rethrown.
  *

packages/nextjs/test/config/withSentry.test.ts

@@ -96,6 +96,10 @@ describe('withSentry', () => {
     });
 
     it('flushes events before finishing non-erroring response', async () => {
+      jest
+        .spyOn(hub.Hub.prototype, 'getClient')
+        .mockReturnValueOnce({ getOptions: () => ({ tracesSampleRate: 1 } as ClientOptions) } as Client);
+
       await callWrappedHandler(wrappedHandlerNoError, req, res);
 
       expect(flushSpy).toHaveBeenCalled();

packages/nextjs/test/config/wrappers.test.ts

@@ -24,7 +24,7 @@ describe('data-fetching function wrappers', () => {
   describe('starts a transaction and puts request in metadata if tracing enabled', () => {
     beforeEach(() => {
       req = { headers: {}, url: 'http://dogs.are.great/tricks/kangaroo' } as IncomingMessage;
-      res = {} as ServerResponse;
+      res = { end: jest.fn() } as unknown as ServerResponse;
 
       jest.spyOn(SentryTracing, 'hasTracingEnabled').mockReturnValueOnce(true);
     });