fix(nextjs): Let flush finish in API routes (#3811)

lobsterkatie · web-flow · commit 75ae8c7c7b5e · 2021-07-20T07:40:30.000-07:00
As discussed in more detail in the penultimate paragraph of #3786, when deployed to vercel, API route lambdas have been shutting down too soon for sentry events to get reliably sent to our servers (in spite of the use of `flush`). This fixes that by wrapping the response's `.end()` method (which is what triggers the lambda shutdown) such that it waits for `flush` to finish before emitting its `finished` signal. Logs from API routes now consistently look like this: ``` [GET] /api/hello Sentry Logger [Log]: [Tracing] Starting http.server transaction - GET /api/hello Sentry Logger [Log]: [Tracing] Finishing http.server transaction - GET /api/hello Sentry Logger [Log]: Flushing events... Sentry Logger [Log]: Done flushing events ```
diff --git a/packages/nextjs/src/utils/handlers.ts b/packages/nextjs/src/utils/handlers.ts
@@ -1,21 +1,27 @@
-import { captureException, flush, getCurrentHub, Handlers, startTransaction, withScope } from '@sentry/node';
-import { extractTraceparentData, getActiveTransaction, hasTracingEnabled } from '@sentry/tracing';
+import { captureException, flush, getCurrentHub, Handlers, startTransaction } from '@sentry/node';
+import { extractTraceparentData, hasTracingEnabled } from '@sentry/tracing';
+import { Transaction } from '@sentry/types';
 import { addExceptionMechanism, isString, logger, stripUrlQueryAndFragment } from '@sentry/utils';
 import * as domain from 'domain';
-import { NextApiHandler } from 'next';
-
-import { addRequestDataToEvent, NextRequest } from './instrumentServer';
+import { NextApiHandler, NextApiResponse } from 'next';
 
 const { parseRequest } = Handlers;
 
 // purely for clarity
 type WrappedNextApiHandler = NextApiHandler;
 
+type AugmentedResponse = NextApiResponse & { __sentryTransaction?: Transaction };
+
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const withSentry = (handler: NextApiHandler): WrappedNextApiHandler => {
   // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
   return async (req, res) => {
-    // wrap everything in a domain in order to prevent scope bleed between requests
+    // 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);
     local.add(res);
@@ -24,73 +30,102 @@ export const withSentry = (handler: NextApiHandler): WrappedNextApiHandler => {
     // 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 () => {
-      try {
-        const currentScope = getCurrentHub().getScope();
+      const currentScope = getCurrentHub().getScope();
 
-        if (currentScope) {
-          currentScope.addEventProcessor(event => addRequestDataToEvent(event, req as NextRequest));
-
-          if (hasTracingEnabled()) {
-            // If there is a trace header set, extract the data from it (parentSpanId, traceId, and sampling decision)
-            let traceparentData;
-            if (req.headers && isString(req.headers['sentry-trace'])) {
-              traceparentData = extractTraceparentData(req.headers['sentry-trace'] as string);
-              logger.log(`[Tracing] Continuing trace ${traceparentData?.traceId}.`);
-            }
+      if (currentScope) {
+        currentScope.addEventProcessor(event => parseRequest(event, req));
+
+        if (hasTracingEnabled()) {
+          // If there is a trace header set, extract the data from it (parentSpanId, traceId, and sampling decision)
+          let traceparentData;
+          if (req.headers && isString(req.headers['sentry-trace'])) {
+            traceparentData = extractTraceparentData(req.headers['sentry-trace'] as string);
+            logger.log(`[Tracing] Continuing trace ${traceparentData?.traceId}.`);
+          }
 
-            const url = `${req.url}`;
-            // pull off query string, if any
-            let reqPath = stripUrlQueryAndFragment(url);
-            // Replace with placeholder
-            if (req.query) {
-              // TODO get this from next if possible, to avoid accidentally replacing non-dynamic parts of the path if
-              // they match dynamic parts
-              for (const [key, value] of Object.entries(req.query)) {
-                reqPath = reqPath.replace(`${value}`, `[${key}]`);
-              }
+          const url = `${req.url}`;
+          // pull off query string, if any
+          let reqPath = stripUrlQueryAndFragment(url);
+          // Replace with placeholder
+          if (req.query) {
+            // TODO get this from next if possible, to avoid accidentally replacing non-dynamic parts of the path if
+            // they match dynamic parts
+            for (const [key, value] of Object.entries(req.query)) {
+              reqPath = reqPath.replace(`${value}`, `[${key}]`);
             }
-            const reqMethod = `${(req.method || 'GET').toUpperCase()} `;
-
-            const transaction = startTransaction(
-              {
-                name: `${reqMethod}${reqPath}`,
-                op: 'http.server',
-                ...traceparentData,
-              },
-              // extra context passed to the `tracesSampler`
-              { request: req },
-            );
-            currentScope.setSpan(transaction);
           }
+          const reqMethod = `${(req.method || 'GET').toUpperCase()} `;
+
+          const transaction = startTransaction(
+            {
+              name: `${reqMethod}${reqPath}`,
+              op: 'http.server',
+              ...traceparentData,
+            },
+            // extra context passed to the `tracesSampler`
+            { request: req },
+          );
+          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 as AugmentedResponse).__sentryTransaction = transaction;
         }
+      }
 
+      try {
         return await handler(req, res); // Call original handler
       } catch (e) {
-        withScope(scope => {
-          scope.addEventProcessor(event => {
+        if (currentScope) {
+          currentScope.addEventProcessor(event => {
             addExceptionMechanism(event, {
               handled: false,
             });
-            return parseRequest(event, req);
+            return event;
           });
           captureException(e);
-        });
-        throw e;
-      } finally {
-        const transaction = getActiveTransaction();
-        if (transaction) {
-          transaction.setHttpStatus(res.statusCode);
-
-          transaction.finish();
-        }
-        try {
-          await flush(2000);
-        } catch (e) {
-          // no-empty
         }
+        throw e;
       }
     });
 
     return await boundHandler();
   };
 };
+
+type ResponseEndMethod = AugmentedResponse['end'];
+type WrappedResponseEndMethod = AugmentedResponse['end'];
+
+function wrapEndMethod(origEnd: ResponseEndMethod): WrappedResponseEndMethod {
+  return async function newEnd(this: AugmentedResponse, ...args: unknown[]) {
+    // TODO: if the handler errored, it will have popped us out of the domain, so all of our scope data will be missing
+
+    const transaction = this.__sentryTransaction;
+
+    if (transaction) {
+      transaction.setHttpStatus(this.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
+    try {
+      logger.log('Flushing events...');
+      await flush(2000);
+      logger.log('Done flushing events');
+    } catch (e) {
+      logger.log(`Error while flushing events:\n${e}`);
+    }
+
+    return origEnd.call(this, ...args);
+  };
+}
diff --git a/packages/nextjs/src/utils/instrumentServer.ts b/packages/nextjs/src/utils/instrumentServer.ts
@@ -1,13 +1,14 @@
-import { captureException, deepReadDirSync, getCurrentHub, startTransaction } from '@sentry/node';
+import { captureException, deepReadDirSync, getCurrentHub, Handlers, startTransaction } from '@sentry/node';
 import { extractTraceparentData, getActiveTransaction, hasTracingEnabled } from '@sentry/tracing';
-import { Event as SentryEvent } from '@sentry/types';
 import { fill, isString, logger, stripUrlQueryAndFragment } from '@sentry/utils';
 import * as domain from 'domain';
 import * as http from 'http';
 import { default as createNextServer } from 'next';
 import * as querystring from 'querystring';
 import * as url from 'url';
 
+const { parseRequest } = Handlers;
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type PlainObject<T = any> = { [key: string]: T };
 
@@ -193,7 +194,7 @@ function makeWrappedReqHandler(origReqHandler: ReqHandler): WrappedReqHandler {
       const currentScope = getCurrentHub().getScope();
 
       if (currentScope) {
-        currentScope.addEventProcessor(event => addRequestDataToEvent(event, req));
+        currentScope.addEventProcessor(event => parseRequest(event, req));
 
         // We only want to record page and API requests
         if (hasTracingEnabled() && shouldTraceRequest(req.url, publicDirFiles)) {
@@ -288,27 +289,3 @@ function shouldTraceRequest(url: string, publicDirFiles: Set<string>): boolean {
   // `static` is a deprecated but still-functional location for static resources
   return !url.startsWith('/_next/') && !url.startsWith('/static/') && !publicDirFiles.has(url);
 }
-
-/**
- * Harvest specific data from the request, and add it to the event.
- *
- * @param event The event to which to add request data
- * @param req The request whose data is being added
- * @returns The modified event
- */
-export function addRequestDataToEvent(event: SentryEvent, req: NextRequest): SentryEvent {
-  // TODO (breaking change): Replace all calls to this function with `parseRequest(event, req, { transaction: false })`
-  // (this is breaking because doing so will change `event.request.url` from a path into an absolute URL, which might
-  // mess up various automations in the Sentry web app - alert rules, auto assignment, saved searches, etc)
-  event.request = {
-    ...event.request,
-    data: req.body,
-    url: req.url.split('?')[0],
-    cookies: req.cookies,
-    headers: req.headers,
-    method: req.method,
-    query_string: req.query,
-  };
-
-  return event;
-}
diff --git a/packages/nextjs/test/integration/test/runner.js b/packages/nextjs/test/integration/test/runner.js
@@ -52,7 +52,11 @@ const runScenario = async (scenario, execute, env) => {
 };
 
 const runScenarios = async (scenarios, execute, env) => {
-  return Promise.all(scenarios.map(scenario => runScenario(scenario, execute, env)));
+  return Promise.all(
+    scenarios.map(scenario => {
+      return runScenario(scenario, execute, env);
+    }),
+  );
 };
 
 module.exports.run = async ({
diff --git a/packages/nextjs/test/integration/test/server/errorApiEndpoint.js b/packages/nextjs/test/integration/test/server/errorApiEndpoint.js
@@ -1,10 +1,12 @@
 const assert = require('assert');
 
 const { sleep } = require('../utils/common');
-const { getAsync, interceptEventRequest } = require('../utils/server');
+const { getAsync, interceptEventRequest, interceptTracingRequest } = require('../utils/server');
 
-module.exports = async ({ url, argv }) => {
-  const capturedRequest = interceptEventRequest(
+module.exports = async ({ url: urlBase, argv }) => {
+  const url = `${urlBase}/api/error`;
+
+  const capturedErrorRequest = interceptEventRequest(
     {
       exception: {
         values: [
@@ -18,7 +20,7 @@ module.exports = async ({ url, argv }) => {
         runtime: 'node',
       },
       request: {
-        url: `${url}/api/error`,
+        url,
         method: 'GET',
       },
       transaction: 'GET /api/error',
@@ -27,8 +29,28 @@ module.exports = async ({ url, argv }) => {
     'errorApiEndpoint',
   );
 
-  await getAsync(`${url}/api/error`);
+  const capturedTransactionRequest = interceptTracingRequest(
+    {
+      contexts: {
+        trace: {
+          op: 'http.server',
+          status: 'internal_error',
+          tags: { 'http.status_code': '500' },
+        },
+      },
+      transaction: 'GET /api/error',
+      type: 'transaction',
+      request: {
+        url,
+      },
+    },
+    argv,
+    'errorApiEndpoint',
+  );
+
+  await getAsync(url);
   await sleep(100);
 
-  assert.ok(capturedRequest.isDone(), 'Did not intercept expected request');
+  assert.ok(capturedErrorRequest.isDone(), 'Did not intercept expected error request');
+  assert.ok(capturedTransactionRequest.isDone(), 'Did not intercept expected transaction request');
 };
diff --git a/packages/nextjs/test/integration/test/server/errorServerSideProps.js b/packages/nextjs/test/integration/test/server/errorServerSideProps.js
@@ -3,7 +3,9 @@ const assert = require('assert');
 const { sleep } = require('../utils/common');
 const { getAsync, interceptEventRequest } = require('../utils/server');
 
-module.exports = async ({ url, argv }) => {
+module.exports = async ({ url: urlBase, argv }) => {
+  const url = `${urlBase}/withServerSideProps`;
+
   const capturedRequest = interceptEventRequest(
     {
       exception: {
@@ -18,15 +20,15 @@ module.exports = async ({ url, argv }) => {
         runtime: 'node',
       },
       request: {
-        url: '/withServerSideProps',
+        url,
         method: 'GET',
       },
     },
     argv,
     'errorServerSideProps',
   );
 
-  await getAsync(`${url}/withServerSideProps`);
+  await getAsync(url);
   await sleep(100);
 
   assert.ok(capturedRequest.isDone(), 'Did not intercept expected request');
diff --git a/packages/nextjs/test/integration/test/server/tracing200.js b/packages/nextjs/test/integration/test/server/tracing200.js
@@ -3,7 +3,9 @@ const assert = require('assert');
 const { sleep } = require('../utils/common');
 const { getAsync, interceptTracingRequest } = require('../utils/server');
 
-module.exports = async ({ url, argv }) => {
+module.exports = async ({ url: urlBase, argv }) => {
+  const url = `${urlBase}/api/users`;
+
   const capturedRequest = interceptTracingRequest(
     {
       contexts: {
@@ -16,14 +18,14 @@ module.exports = async ({ url, argv }) => {
       transaction: 'GET /api/users',
       type: 'transaction',
       request: {
-        url: '/api/users',
+        url,
       },
     },
     argv,
     'tracing200',
   );
 
-  await getAsync(`${url}/api/users`);
+  await getAsync(url);
   await sleep(100);
 
   assert.ok(capturedRequest.isDone(), 'Did not intercept expected request');
diff --git a/packages/nextjs/test/integration/test/server/tracing500.js b/packages/nextjs/test/integration/test/server/tracing500.js
@@ -3,7 +3,8 @@ const assert = require('assert');
 const { sleep } = require('../utils/common');
 const { getAsync, interceptTracingRequest } = require('../utils/server');
 
-module.exports = async ({ url, argv }) => {
+module.exports = async ({ url: urlBase, argv }) => {
+  const url = `${urlBase}/api/broken`;
   const capturedRequest = interceptTracingRequest(
     {
       contexts: {
@@ -16,14 +17,14 @@ module.exports = async ({ url, argv }) => {
       transaction: 'GET /api/broken',
       type: 'transaction',
       request: {
-        url: '/api/broken',
+        url,
       },
     },
     argv,
     'tracing500',
   );
 
-  await getAsync(`${url}/api/broken`);
+  await getAsync(url);
   await sleep(100);
 
   assert.ok(capturedRequest.isDone(), 'Did not intercept expected request');
diff --git a/packages/nextjs/test/integration/test/server/tracingHttp.js b/packages/nextjs/test/integration/test/server/tracingHttp.js
