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

When deployed to Vercel, nextjs apps have their route handlers converted to AWS lambdas, which use the vercel node runtime. In this runtime, AWS's [`callbackWaitsForEmptyEventLoop` context option](https://docs.aws.amazon.com/lambda/latest/dg/nodejs-context.html) is [set to `false`](https://github.com/vercel/fun/blob/6329dcbd3f7d99a519cbbd6ac615fec4b46fc28e/src/runtimes/nodejs/bootstrap.ts#L122). As a result, it's possible for a route-handling lambda to be shut down before events have finished sending to Sentry. To ensure that this doesn't happen, in `withSentry` we wrap `res.end`, so that it flushes all of our events before the response marks itself as finished.

This extracts that logic into its own module, so that it can also be used in the other `withSentryX` wrappers (`withSentryServerSideProps` and the like).

Author: lobsterkatie

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