move newEnd into a factory function

Author: lobsterkatie

packages/nextjs/src/utils/handlers.ts

@@ -1,6 +1,8 @@
 import { captureException, flush, getCurrentHub, Handlers, startTransaction, withScope } from '@sentry/node';
+// import { extractTraceparentData, hasTracingEnabled } from '@sentry/tracing';
 import { extractTraceparentData, getActiveTransaction, hasTracingEnabled } from '@sentry/tracing';
 import { addExceptionMechanism, isString, logger, stripUrlQueryAndFragment } from '@sentry/utils';
+import { Scope } from '@sentry/hub';
 import * as domain from 'domain';
 import { NextApiHandler, NextApiResponse } from 'next';
 
@@ -11,43 +13,23 @@ const { parseRequest } = Handlers;
 // purely for clarity
 type WrappedNextApiHandler = NextApiHandler;
 
+type ScopedResponse = NextApiResponse & { __sentryScope?: Scope };
+
 // 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) => {
-    const origEnd = res.end;
-
-    async function newEnd(this: NextApiResponse, ...args: any[]) {
-      const transaction = getActiveTransaction();
-      console.log('Active transaction:', transaction?.name);
-
-      if (transaction) {
-        transaction.setHttpStatus(res.statusCode);
-
-        transaction.finish();
-      }
-      try {
-        logger.log('Flushing events...');
-        await flush(2000);
-      } catch (e) {
-        logger.log(`Error while flushing events:\n${e}`);
-      } finally {
-        logger.log('Done flushing events');
-        // res.end();
-      }
-
-      return origEnd.call(this, ...args);
-    }
+    // 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)
+    res.end = wrapEndMethod(res.end);
 
-    res.end = newEnd;
-
-    // wrap everything in a domain in order to prevent scope bleed between requests
+    // use a domain in order to prevent scope bleed between requests
     const local = domain.create();
     local.add(req);
     local.add(res);
 
     return local.bind(async () => {
-      try {
+    
         const currentScope = getCurrentHub().getScope();
 
         if (currentScope) {
@@ -84,11 +66,16 @@ export const withSentry = (handler: NextApiHandler): WrappedNextApiHandler => {
               { 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 the transaction and attach the correct data to it
+            (res as ScopedResponse).__sentryScope = currentScope;
           }
         }
-
+        try {
         return await handler(req, res); // Call original handler
       } catch (e) {
+        // TODO how do we know
         withScope(scope => {
           scope.addEventProcessor(event => {
             addExceptionMechanism(event, {
@@ -119,3 +106,48 @@ export const withSentry = (handler: NextApiHandler): WrappedNextApiHandler => {
     })();
   };
 };
+
+type ResponseEndMethod = ScopedResponse['end'];
+type WrappedResponseEndMethod = ScopedResponse['end'];
+
+function wrapEndMethod(origEnd: ResponseEndMethod): WrappedResponseEndMethod {
+  return async function newEnd(this: ScopedResponse, ...args: any[]) {
+    // if the handler errored, it will have popped us out of the domain, so push the domain's scope onto the stack
+    // just in case (if we *are* still in the domain, this will replace the current scope with a clone of itself,
+    // which is effectively a no-op as long as we remember to pop it off when we're done)
+    const currentHub = getCurrentHub();
+    currentHub.pushScope(this.__sentryScope);
+
+    const transaction = getActiveTransaction();
+    console.log('Active transaction:', transaction?.name);
+
+    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);
+    } catch (e) {
+      logger.log(`Error while flushing events:\n${e}`);
+    } finally {
+      logger.log('Done flushing events');
+    }
+
+    // now that our work is done, we can pop off the scope and allow the response to end
+    currentHub.popScope();
+    return origEnd.call(this, ...args);
+  };
+}