fix(remix): Provide sentry-trace and baggage via root loader.

Author: onurtemizkan

packages/remix/src/utils/instrumentServer.ts

@@ -1,9 +1,10 @@
 /* eslint-disable max-lines */
 import { captureException, getCurrentHub } from '@sentry/node';
-import { getActiveTransaction } from '@sentry/tracing';
+import { getActiveTransaction, hasTracingEnabled } from '@sentry/tracing';
 import {
   addExceptionMechanism,
   fill,
+  isNodeEnv,
   loadModule,
   logger,
   serializeBaggage,
@@ -79,6 +80,8 @@ interface HandleDataRequestFunction {
 
 interface ServerEntryModule {
   default: HandleDocumentRequestFunction;
+  meta: MetaFunction;
+  loader: DataFunction;
   handleDataRequest?: HandleDataRequestFunction;
 }
 
@@ -232,33 +235,30 @@ function makeWrappedLoader(origAction: DataFunction): DataFunction {
   return makeWrappedDataFunction(origAction, 'loader');
 }
 
-function makeWrappedMeta(origMeta: MetaFunction | HtmlMetaDescriptor = {}): MetaFunction {
-  return function (
-    this: unknown,
-    args: { data: AppData; parentsData: RouteData; params: Params; location: Location },
-  ): HtmlMetaDescriptor {
-    let origMetaResult;
-    if (origMeta instanceof Function) {
-      origMetaResult = origMeta.call(this, args);
-    } else {
-      origMetaResult = origMeta;
-    }
+function makeWrappedRootLoader(origLoader: DataFunction): DataFunction {
+  return async function (this: unknown, args: DataFunctionArgs): Promise<Response | AppData> {
+    let sentryTrace;
+    let sentryBaggage;
+
+    const activeTransaction = getActiveTransaction();
+    const currentScope = getCurrentHub().getScope();
+
+    if (isNodeEnv() && hasTracingEnabled()) {
+      if (currentScope) {
+        const span = currentScope.getSpan();
 
-    const scope = getCurrentHub().getScope();
-    if (scope) {
-      const span = scope.getSpan();
-      const transaction = getActiveTransaction();
-
-      if (span && transaction) {
-        return {
-          ...origMetaResult,
-          'sentry-trace': span.toTraceparent(),
-          baggage: serializeBaggage(transaction.getBaggage()),
-        };
+        if (span && activeTransaction) {
+          sentryTrace = span.toTraceparent();
+          sentryBaggage = serializeBaggage(activeTransaction.getBaggage());
+        }
       }
     }
 
-    return origMetaResult;
+    const res = await origLoader.call(this, args);
+
+    Object.assign(res, { sentryTrace, sentryBaggage });
+
+    return res;
   };
 }
 
@@ -303,12 +303,20 @@ function makeWrappedCreateRequestHandler(
     for (const [id, route] of Object.entries(build.routes)) {
       const wrappedRoute = { ...route, module: { ...route.module } };
 
-      fill(wrappedRoute.module, 'meta', makeWrappedMeta);
-
       if (wrappedRoute.module.action) {
         fill(wrappedRoute.module, 'action', makeWrappedAction);
       }
 
+      // Entry module should have a loader function to provide `sentry-trace` and `baggage`
+      // They will be available for the root `meta` function as `data.sentryTrace` and `data.sentryBaggage`
+      if (!wrappedRoute.parentId) {
+        if (!wrappedRoute.module.loader) {
+          wrappedRoute.module.loader = () => ({});
+        }
+
+        fill(wrappedRoute.module, 'loader', makeWrappedRootLoader);
+      }
+
       if (wrappedRoute.module.loader) {
         fill(wrappedRoute.module, 'loader', makeWrappedLoader);
       }