feat(nextjs): Use Vercel's waitUntil to defer freezing of Vercel Lambdas (#12133)

Luca Forstner · mydea · web-flow · commit 2043f2d575ab · 2024-05-27T11:48:56.000+02:00
Co-authored-by: Francesco Novy &lt;francesco.novy@sentry.io&gt;
diff --git a/packages/nextjs/src/common/_error.ts b/packages/nextjs/src/common/_error.ts
@@ -1,6 +1,7 @@
 import { captureException, withScope } from '@sentry/core';
 import type { NextPageContext } from 'next';
-import { flushQueue } from './utils/responseEnd';
+import { flushSafelyWithTimeout } from './utils/responseEnd';
+import { vercelWaitUntil } from './utils/vercelWaitUntil';
 
 type ContextOrProps = {
   req?: NextPageContext['req'];
@@ -53,7 +54,5 @@ export async function captureUnderscoreErrorException(contextOrProps: ContextOrP
     });
   });
 
-  // In case this is being run as part of a serverless function (as is the case with the server half of nextjs apps
-  // deployed to vercel), make sure the error gets sent to Sentry before the lambda exits.
-  await flushQueue();
+  vercelWaitUntil(flushSafelyWithTimeout());
 }
diff --git a/packages/nextjs/src/common/types.ts b/packages/nextjs/src/common/types.ts
@@ -48,11 +48,6 @@ export type VercelCronsConfig = { path?: string; schedule?: string }[] | undefin
 export type NextApiHandler = {
   (req: NextApiRequest, res: NextApiResponse): void | Promise<void> | unknown | Promise<unknown>;
   __sentry_route__?: string;
-
-  /**
-   * A property we set in our integration tests to simulate running an API route on platforms that don't support streaming.
-   */
-  __sentry_test_doesnt_support_streaming__?: true;
 };
 
 export type WrappedNextApiHandler = {
diff --git a/packages/nextjs/src/common/utils/edgeWrapperUtils.ts b/packages/nextjs/src/common/utils/edgeWrapperUtils.ts
@@ -12,8 +12,9 @@ import {
 import { winterCGRequestToRequestData } from '@sentry/utils';
 
 import type { EdgeRouteHandler } from '../../edge/types';
-import { flushQueue } from './responseEnd';
+import { flushSafelyWithTimeout } from './responseEnd';
 import { commonObjectToIsolationScope, escapeNextjsTracing } from './tracingUtils';
+import { vercelWaitUntil } from './vercelWaitUntil';
 
 /**
  * Wraps a function on the edge runtime with error and performance monitoring.
@@ -80,9 +81,11 @@ export function withEdgeWrapping<H extends EdgeRouteHandler>(
 
                 return handlerResult;
               },
-            ).finally(() => flushQueue());
+            );
           },
-        );
+        ).finally(() => {
+          vercelWaitUntil(flushSafelyWithTimeout());
+        });
       });
     });
   };
diff --git a/packages/nextjs/src/common/utils/platformSupportsStreaming.ts b/packages/nextjs/src/common/utils/platformSupportsStreaming.ts
diff --git a/packages/nextjs/src/common/utils/responseEnd.ts b/packages/nextjs/src/common/utils/responseEnd.ts
@@ -44,10 +44,14 @@ export function finishSpan(span: Span, res: ServerResponse): void {
   span.end();
 }
 
-/** 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> {
+/**
+ * Flushes pending Sentry events with a 2 second timeout and in a way that cannot create unhandled promise rejections.
+ */
+export async function flushSafelyWithTimeout(): Promise<void> {
   try {
     DEBUG_BUILD && logger.log('Flushing events...');
+    // We give things that are currently stuck in event processors a tiny bit more time to finish before flushing. 50ms was chosen very unscientifically.
+    await new Promise(resolve => setTimeout(resolve, 50));
     await flush(2000);
     DEBUG_BUILD && logger.log('Done flushing events');
   } catch (e) {
diff --git a/packages/nextjs/src/common/utils/vercelWaitUntil.ts b/packages/nextjs/src/common/utils/vercelWaitUntil.ts
@@ -0,0 +1,21 @@
+import { GLOBAL_OBJ } from '@sentry/utils';
+
+interface VercelRequestContextGlobal {
+  get?(): {
+    waitUntil?: (task: Promise<unknown>) => void;
+  };
+}
+
+/**
+ * Function that delays closing of a Vercel lambda until the provided promise is resolved.
+ *
+ * Vendored from https://www.npmjs.com/package/@vercel/functions
+ */
+export function vercelWaitUntil(task: Promise<unknown>): void {
+  const vercelRequestContextGlobal: VercelRequestContextGlobal | undefined =
+    // @ts-expect-error This is not typed
+    GLOBAL_OBJ[Symbol.for('@vercel/request-context')];
+
+  const ctx = vercelRequestContextGlobal?.get?.() ?? {};
+  ctx.waitUntil?.(task);
+}
diff --git a/packages/nextjs/src/common/utils/wrapperUtils.ts b/packages/nextjs/src/common/utils/wrapperUtils.ts
@@ -15,9 +15,9 @@ import {
 import type { Span } from '@sentry/types';
 import { isString } from '@sentry/utils';
 
-import { platformSupportsStreaming } from './platformSupportsStreaming';
-import { autoEndSpanOnResponseEnd, flushQueue } from './responseEnd';
+import { autoEndSpanOnResponseEnd, flushSafelyWithTimeout } from './responseEnd';
 import { commonObjectToIsolationScope, escapeNextjsTracing } from './tracingUtils';
+import { vercelWaitUntil } from './vercelWaitUntil';
 
 declare module 'http' {
   interface IncomingMessage {
@@ -124,15 +124,14 @@ export function withTracedServerSideDataFetcher<F extends (...args: any[]) => Pr
                   throw e;
                 } finally {
                   dataFetcherSpan.end();
-                  if (!platformSupportsStreaming()) {
-                    await flushQueue();
-                  }
                 }
               },
             );
           });
         });
       });
+    }).finally(() => {
+      vercelWaitUntil(flushSafelyWithTimeout());
     });
   };
 }
@@ -198,10 +197,9 @@ export async function callDataFetcherTraced<F extends (...args: any[]) => Promis
         throw e;
       } finally {
         dataFetcherSpan.end();
-        if (!platformSupportsStreaming()) {
-          await flushQueue();
-        }
       }
     },
-  );
+  ).finally(() => {
+    vercelWaitUntil(flushSafelyWithTimeout());
+  });
 }
diff --git a/packages/nextjs/src/common/withServerActionInstrumentation.ts b/packages/nextjs/src/common/withServerActionInstrumentation.ts
@@ -9,9 +9,9 @@ import { logger } from '@sentry/utils';
 
 import { DEBUG_BUILD } from './debug-build';
 import { isNotFoundNavigationError, isRedirectNavigationError } from './nextNavigationErrorUtils';
-import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
-import { flushQueue } from './utils/responseEnd';
+import { flushSafelyWithTimeout } from './utils/responseEnd';
 import { escapeNextjsTracing } from './utils/tracingUtils';
+import { vercelWaitUntil } from './utils/vercelWaitUntil';
 
 interface Options {
   formData?: FormData;
@@ -131,16 +131,7 @@ async function withServerActionInstrumentationImplementation<A extends (...args:
               },
             );
           } finally {
-            if (!platformSupportsStreaming()) {
-              // Lambdas require manual flushing to prevent execution freeze before the event is sent
-              await flushQueue();
-            }
-
-            if (process.env.NEXT_RUNTIME === 'edge') {
-              // flushQueue should not throw
-              // eslint-disable-next-line @typescript-eslint/no-floating-promises
-              flushQueue();
-            }
+            vercelWaitUntil(flushSafelyWithTimeout());
           }
         },
       );
diff --git a/packages/nextjs/src/common/wrapApiHandlerWithSentry.ts b/packages/nextjs/src/common/wrapApiHandlerWithSentry.ts
@@ -10,9 +10,9 @@ import { consoleSandbox, isString, logger, objectify } from '@sentry/utils';
 
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
 import type { AugmentedNextApiRequest, AugmentedNextApiResponse, NextApiHandler } from './types';
-import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
-import { flushQueue } from './utils/responseEnd';
+import { flushSafelyWithTimeout } from './utils/responseEnd';
 import { escapeNextjsTracing } from './utils/tracingUtils';
+import { vercelWaitUntil } from './utils/vercelWaitUntil';
 
 /**
  * Wrap the given API route handler for tracing and error capturing. Thin wrapper around `withSentry`, which only
@@ -83,15 +83,8 @@ export function wrapApiHandlerWithSentry(apiHandler: NextApiHandler, parameteriz
                     apply(target, thisArg, argArray) {
                       setHttpStatus(span, res.statusCode);
                       span.end();
-                      if (platformSupportsStreaming() && !wrappingTarget.__sentry_test_doesnt_support_streaming__) {
-                        target.apply(thisArg, argArray);
-                      } else {
-                        // flushQueue will not reject
-                        // eslint-disable-next-line @typescript-eslint/no-floating-promises
-                        flushQueue().then(() => {
-                          target.apply(thisArg, argArray);
-                        });
-                      }
+                      vercelWaitUntil(flushSafelyWithTimeout());
+                      target.apply(thisArg, argArray);
                     },
                   });
 
@@ -138,14 +131,7 @@ export function wrapApiHandlerWithSentry(apiHandler: NextApiHandler, parameteriz
                     setHttpStatus(span, res.statusCode);
                     span.end();
 
-                    // 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 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.)
-                    if (platformSupportsStreaming() && !wrappingTarget.__sentry_test_doesnt_support_streaming__) {
-                      await flushQueue();
-                    }
+                    vercelWaitUntil(flushSafelyWithTimeout());
 
                     // 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/common/wrapRouteHandlerWithSentry.ts b/packages/nextjs/src/common/wrapRouteHandlerWithSentry.ts
@@ -13,13 +13,13 @@ import {
 import { propagationContextFromHeaders, winterCGHeadersToDict } from '@sentry/utils';
 import { isNotFoundNavigationError, isRedirectNavigationError } from './nextNavigationErrorUtils';
 import type { RouteHandlerContext } from './types';
-import { platformSupportsStreaming } from './utils/platformSupportsStreaming';
-import { flushQueue } from './utils/responseEnd';
+import { flushSafelyWithTimeout } from './utils/responseEnd';
 import {
   commonObjectToIsolationScope,
   commonObjectToPropagationContext,
   escapeNextjsTracing,
 } from './utils/tracingUtils';
+import { vercelWaitUntil } from './utils/vercelWaitUntil';
 
 /**
  * Wraps a Next.js route handler with performance and error instrumentation.
@@ -97,11 +97,7 @@ export function wrapRouteHandlerWithSentry<F extends (...args: any[]) => any>(
                 },
               );
             } finally {
-              if (!platformSupportsStreaming() || process.env.NEXT_RUNTIME === 'edge') {
-                // 1. Edge transport requires manual flushing
-                // 2. Lambdas require manual flushing to prevent execution freeze before the event is sent
-                await flushQueue();
-              }
+              vercelWaitUntil(flushSafelyWithTimeout());
             }
           });
         });
diff --git a/packages/nextjs/src/common/wrapServerComponentWithSentry.ts b/packages/nextjs/src/common/wrapServerComponentWithSentry.ts
@@ -14,12 +14,13 @@ import { propagationContextFromHeaders, uuid4, winterCGHeadersToDict } from '@se
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
 import { isNotFoundNavigationError, isRedirectNavigationError } from '../common/nextNavigationErrorUtils';
 import type { ServerComponentContext } from '../common/types';
-import { flushQueue } from './utils/responseEnd';
+import { flushSafelyWithTimeout } from './utils/responseEnd';
 import {
   commonObjectToIsolationScope,
   commonObjectToPropagationContext,
   escapeNextjsTracing,
 } from './utils/tracingUtils';
+import { vercelWaitUntil } from './utils/vercelWaitUntil';
 
 /**
  * Wraps an `app` directory server component with Sentry error instrumentation.
@@ -93,10 +94,7 @@ export function wrapServerComponentWithSentry<F extends (...args: any[]) => any>
                   },
                   () => {
                     span.end();
-
-                    // flushQueue should not throw
-                    // eslint-disable-next-line @typescript-eslint/no-floating-promises
-                    flushQueue();
+                    vercelWaitUntil(flushSafelyWithTimeout());
                   },
                 );
               },
diff --git a/packages/nextjs/test/integration/pages/api/doubleEndMethodOnVercel.ts b/packages/nextjs/test/integration/pages/api/doubleEndMethodOnVercel.ts
@@ -6,6 +6,4 @@ const handler = async (_req: NextApiRequest, res: NextApiResponse): Promise<void
   res.end();
 };
 
-handler.__sentry_test_doesnt_support_streaming__ = true;
-
 export default handler;
