feat(nextjs): Send events consistently on platforms that don't support streaming (#6578)

Luca Forstner · web-flow · commit 371d42d6eb83 · 2023-01-03T17:16:08.000+01:00
diff --git a/packages/nextjs/src/config/wrappers/utils/responseEnd.ts b/packages/nextjs/src/config/wrappers/utils/responseEnd.ts
@@ -24,10 +24,8 @@ import { ResponseEndMethod, WrappedResponseEndMethod } from '../types';
  */
 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 function sentryWrappedEnd(this: ServerResponse, ...args: unknown[]) {
+      void finishTransaction(transaction, this);
       return origEnd.call(this, ...args);
     };
   };
diff --git a/packages/nextjs/src/config/wrappers/withSentryAPI.ts b/packages/nextjs/src/config/wrappers/withSentryAPI.ts
@@ -1,5 +1,6 @@
 import { captureException, getCurrentHub, startTransaction } from '@sentry/node';
 import { extractTraceparentData, hasTracingEnabled } from '@sentry/tracing';
+import { Transaction } from '@sentry/types';
 import {
   addExceptionMechanism,
   baggageHeaderToDynamicSamplingContext,
@@ -11,6 +12,7 @@ import {
 import * as domain from 'domain';
 
 import { formatAsCode, nextLogger } from '../../utils/nextLogger';
+import { platformSupportsStreaming } from '../../utils/platformSupportsStreaming';
 import type { AugmentedNextApiRequest, AugmentedNextApiResponse, NextApiHandler, WrappedNextApiHandler } from './types';
 import { autoEndTransactionOnResponseEnd, finishTransaction, flushQueue } from './utils/responseEnd';
 
@@ -74,8 +76,9 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
     // `local.bind` causes everything to run inside a domain, just like `local.run` does, but it also lets the callback
     // 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.
+    // eslint-disable-next-line complexity
     const boundHandler = local.bind(async () => {
-      let transaction;
+      let transaction: Transaction | undefined;
       const currentScope = getCurrentHub().getScope();
 
       if (currentScope) {
@@ -127,7 +130,43 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
           );
           currentScope.setSpan(transaction);
 
-          autoEndTransactionOnResponseEnd(transaction, res);
+          if (platformSupportsStreaming()) {
+            autoEndTransactionOnResponseEnd(transaction, res);
+          } else {
+            // If we're not on a platform that supports streaming, we're blocking all response-ending methods until the
+            // queue is flushed.
+
+            const origResSend = res.send;
+            res.send = async function (this: unknown, ...args: unknown[]) {
+              if (transaction) {
+                await finishTransaction(transaction, res);
+                await flushQueue();
+              }
+
+              origResSend.apply(this, args);
+            };
+
+            const origResJson = res.json;
+            res.json = async function (this: unknown, ...args: unknown[]) {
+              if (transaction) {
+                await finishTransaction(transaction, res);
+                await flushQueue();
+              }
+
+              origResJson.apply(this, args);
+            };
+
+            // eslint-disable-next-line @typescript-eslint/unbound-method
+            const origResEnd = res.end;
+            res.end = async function (this: unknown, ...args: unknown[]) {
+              if (transaction) {
+                await finishTransaction(transaction, res);
+                await flushQueue();
+              }
+
+              origResEnd.apply(this, args);
+            };
+          }
         }
       }
 
@@ -184,8 +223,12 @@ export function withSentry(origHandler: NextApiHandler, parameterizedRoute?: str
         // 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 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();
+        if (platformSupportsStreaming()) {
+          void finishTransaction(transaction, res);
+        } else {
+          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
diff --git a/packages/nextjs/src/config/wrappers/wrapperUtils.ts b/packages/nextjs/src/config/wrappers/wrapperUtils.ts
@@ -5,7 +5,8 @@ import { baggageHeaderToDynamicSamplingContext, extractTraceparentData } from '@
 import * as domain from 'domain';
 import { IncomingMessage, ServerResponse } from 'http';
 
-import { autoEndTransactionOnResponseEnd } from './utils/responseEnd';
+import { platformSupportsStreaming } from '../../utils/platformSupportsStreaming';
+import { autoEndTransactionOnResponseEnd, flushQueue } from './utils/responseEnd';
 
 declare module 'http' {
   interface IncomingMessage {
@@ -77,43 +78,62 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
   return async function (this: unknown, ...args: Parameters<F>): Promise<ReturnType<F>> {
     return domain.create().bind(async () => {
       let requestTransaction: Transaction | undefined = getTransactionFromRequest(req);
-
-      if (requestTransaction === undefined) {
-        const sentryTraceHeader = req.headers['sentry-trace'];
-        const rawBaggageString = req.headers && req.headers.baggage;
-        const traceparentData =
-          typeof sentryTraceHeader === 'string' ? extractTraceparentData(sentryTraceHeader) : undefined;
-
-        const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(rawBaggageString);
-
-        const newTransaction = startTransaction(
-          {
-            op: 'http.server',
-            name: options.requestedRouteName,
-            ...traceparentData,
-            status: 'ok',
-            metadata: {
-              source: 'route',
-              dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+      let dataFetcherSpan;
+
+      const sentryTraceHeader = req.headers['sentry-trace'];
+      const rawBaggageString = req.headers && req.headers.baggage;
+      const traceparentData =
+        typeof sentryTraceHeader === 'string' ? extractTraceparentData(sentryTraceHeader) : undefined;
+
+      const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(rawBaggageString);
+
+      if (platformSupportsStreaming()) {
+        if (requestTransaction === undefined) {
+          const newTransaction = startTransaction(
+            {
+              op: 'http.server',
+              name: options.requestedRouteName,
+              ...traceparentData,
+              status: 'ok',
+              metadata: {
+                request: req,
+                source: 'route',
+                dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+              },
             },
-          },
-          { request: req },
-        );
+            { request: req },
+          );
 
-        requestTransaction = newTransaction;
-        autoEndTransactionOnResponseEnd(newTransaction, res);
+          requestTransaction = newTransaction;
 
-        // Link the transaction and the request together, so that when we would normally only have access to one, it's
-        // still possible to grab the other.
-        setTransactionOnRequest(newTransaction, req);
-        newTransaction.setMetadata({ request: req });
-      }
+          if (platformSupportsStreaming()) {
+            // On platforms that don't support streaming, doing things after res.end() is unreliable.
+            autoEndTransactionOnResponseEnd(newTransaction, res);
+          }
 
-      const dataFetcherSpan = requestTransaction.startChild({
-        op: 'function.nextjs',
-        description: `${options.dataFetchingMethodName} (${options.dataFetcherRouteName})`,
-        status: 'ok',
-      });
+          // Link the transaction and the request together, so that when we would normally only have access to one, it's
+          // still possible to grab the other.
+          setTransactionOnRequest(newTransaction, req);
+        }
+
+        dataFetcherSpan = requestTransaction.startChild({
+          op: 'function.nextjs',
+          description: `${options.dataFetchingMethodName} (${options.dataFetcherRouteName})`,
+          status: 'ok',
+        });
+      } else {
+        dataFetcherSpan = startTransaction({
+          op: 'function.nextjs',
+          name: `${options.dataFetchingMethodName} (${options.dataFetcherRouteName})`,
+          ...traceparentData,
+          status: 'ok',
+          metadata: {
+            request: req,
+            source: 'route',
+            dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+          },
+        });
+      }
 
       const currentScope = getCurrentHub().getScope();
       if (currentScope) {
@@ -127,15 +147,13 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
         // Since we finish the span before the error can bubble up and trigger the handlers in `registerErrorInstrumentation`
         // that set the transaction status, we need to manually set the status of the span & transaction
         dataFetcherSpan.setStatus('internal_error');
-
-        const transaction = dataFetcherSpan.transaction;
-        if (transaction) {
-          transaction.setStatus('internal_error');
-        }
-
+        requestTransaction?.setStatus('internal_error');
         throw e;
       } finally {
         dataFetcherSpan.finish();
+        if (!platformSupportsStreaming()) {
+          await flushQueue();
+        }
       }
     })();
   };
diff --git a/packages/nextjs/src/utils/platformSupportsStreaming.ts b/packages/nextjs/src/utils/platformSupportsStreaming.ts
@@ -0,0 +1 @@
+export const platformSupportsStreaming = (): boolean => !process.env.LAMBDA_TASK_ROOT && !process.env.VERCEL;
diff --git a/packages/nextjs/test/config/withSentry.test.ts b/packages/nextjs/test/config/withSentry.test.ts
@@ -1,7 +1,6 @@
 import * as hub from '@sentry/core';
 import * as Sentry from '@sentry/node';
 import { Client, ClientOptions } from '@sentry/types';
-import * as utils from '@sentry/utils';
 import { NextApiHandler, NextApiRequest, NextApiResponse } from 'next';
 
 import { AugmentedNextApiResponse, withSentry, WrappedNextApiHandler } from '../../src/config/wrappers';
@@ -36,31 +35,16 @@ async function callWrappedHandler(wrappedHandler: WrappedNextApiHandler, req: Ne
   }
 }
 
-// We mock `captureException` as a no-op because under normal circumstances it is an un-awaited effectively-async
-// function which might or might not finish before any given test ends, potentially leading jest to error out.
-const captureExceptionSpy = jest.spyOn(Sentry, 'captureException').mockImplementation(jest.fn());
-const loggerSpy = jest.spyOn(utils.logger, 'log');
-const flushSpy = jest.spyOn(Sentry, 'flush').mockImplementation(async () => {
-  // simulate the time it takes time to flush all events
-  await sleep(FLUSH_DURATION);
-  return true;
-});
 const startTransactionSpy = jest.spyOn(Sentry, 'startTransaction');
 
 describe('withSentry', () => {
   let req: NextApiRequest, res: NextApiResponse;
 
-  const noShoesError = new Error('Oh, no! Charlie ate the flip-flops! :-(');
-
   const origHandlerNoError: NextApiHandler = async (_req, res) => {
     res.send('Good dog, Maisey!');
   };
-  const origHandlerWithError: NextApiHandler = async (_req, _res) => {
-    throw noShoesError;
-  };
 
   const wrappedHandlerNoError = withSentry(origHandlerNoError);
-  const wrappedHandlerWithError = withSentry(origHandlerWithError);
 
   beforeEach(() => {
     req = { url: 'http://dogs.are.great' } as NextApiRequest;
@@ -78,35 +62,6 @@ describe('withSentry', () => {
     jest.clearAllMocks();
   });
 
-  describe('flushing', () => {
-    it('flushes events before rethrowing error', async () => {
-      try {
-        await callWrappedHandler(wrappedHandlerWithError, req, res);
-      } catch (err) {
-        expect(err).toBe(noShoesError);
-      }
-
-      expect(captureExceptionSpy).toHaveBeenCalledWith(noShoesError);
-      expect(flushSpy).toHaveBeenCalled();
-      expect(loggerSpy).toHaveBeenCalledWith('Done flushing events');
-
-      // This ensures the expect inside the `catch` block actually ran, i.e., that in the end the wrapped handler
-      // errored out the same way it would without sentry, meaning the error was indeed rethrown
-      expect.assertions(4);
-    });
-
-    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();
-      expect(loggerSpy).toHaveBeenCalledWith('Done flushing events');
-    });
-  });
-
   describe('tracing', () => {
     it('starts a transaction and sets metadata when tracing is enabled', async () => {
       jest
diff --git a/packages/nextjs/test/config/wrappers.test.ts b/packages/nextjs/test/config/wrappers.test.ts
@@ -14,7 +14,6 @@ import {
 } from '../../src/config/wrappers';
 
 const startTransactionSpy = jest.spyOn(SentryCore, 'startTransaction');
-const setMetadataSpy = jest.spyOn(SentryTracing.Transaction.prototype, 'setMetadata');
 
 describe('data-fetching function wrappers', () => {
   const route = '/tricks/[trickName]';
@@ -43,16 +42,14 @@ describe('data-fetching function wrappers', () => {
         expect.objectContaining({
           name: '/tricks/[trickName]',
           op: 'http.server',
-          metadata: expect.objectContaining({ source: 'route' }),
+          metadata: expect.objectContaining({ source: 'route', request: req }),
         }),
         {
           request: expect.objectContaining({
             url: 'http://dogs.are.great/tricks/kangaroo',
           }),
         },
       );
-
-      expect(setMetadataSpy).toHaveBeenCalledWith({ request: req });
     });
 
     test('withSentryServerSideGetInitialProps', async () => {
@@ -65,16 +62,14 @@ describe('data-fetching function wrappers', () => {
         expect.objectContaining({
           name: '/tricks/[trickName]',
           op: 'http.server',
-          metadata: expect.objectContaining({ source: 'route' }),
+          metadata: expect.objectContaining({ source: 'route', request: req }),
         }),
         {
           request: expect.objectContaining({
             url: 'http://dogs.are.great/tricks/kangaroo',
           }),
         },
       );
-
-      expect(setMetadataSpy).toHaveBeenCalledWith({ request: req });
     });
   });
 });

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+export const platformSupportsStreaming = (): boolean => !process.env.LAMBDA_TASK_ROOT && !process.env.VERCEL;`