fix(sveltekit): Handle server-only and shared load functions

packages/sveltekit/src/server/load.ts

@@ -1,8 +1,14 @@
 /* eslint-disable @sentry-internal/sdk/no-optional-chaining */
+import { trace } from '@sentry/core';
 import { captureException } from '@sentry/node';
-import { addExceptionMechanism, isThenable, objectify } from '@sentry/utils';
-import type { HttpError, Load, ServerLoad } from '@sveltejs/kit';
-import * as domain from 'domain';
+import type { DynamicSamplingContext, TraceparentData, TransactionContext } from '@sentry/types';
+import {
+  addExceptionMechanism,
+  baggageHeaderToDynamicSamplingContext,
+  extractTraceparentData,
+  objectify,
+} from '@sentry/utils';
+import type { HttpError, Load, LoadEvent, ServerLoad, ServerLoadEvent } from '@sveltejs/kit';
 
 function isHttpError(err: unknown): err is HttpError {
   return typeof err === 'object' && err !== null && 'status' in err && 'body' in err;
@@ -45,25 +51,52 @@ function sendErrorToSentry(e: unknown): unknown {
  */
 export function wrapLoadWithSentry<T extends ServerLoad | Load>(origLoad: T): T {
   return new Proxy(origLoad, {
-    apply: (wrappingTarget, thisArg, args: Parameters<ServerLoad>) => {
-      return domain.create().bind(() => {
-        let maybePromiseResult: ReturnType<T>;
+    apply: (wrappingTarget, thisArg, args: Parameters<ServerLoad | Load>) => {
+      const [event] = args;
+      const routeId = event.route && event.route.id;
 
-        try {
-          maybePromiseResult = wrappingTarget.apply(thisArg, args);
-        } catch (e) {
-          sendErrorToSentry(e);
-          throw e;
-        }
+      const { traceparentData, dynamicSamplingContext } = getTracePropagationData(event);
 
-        if (isThenable(maybePromiseResult)) {
-          Promise.resolve(maybePromiseResult).then(null, e => {
-            sendErrorToSentry(e);
-          });
-        }
+      const traceLoadContext: TransactionContext = {
+        op: 'function.sveltekit.load',
+        name: routeId ? routeId : event.url.pathname,
+        status: 'ok',
+        metadata: {
+          source: routeId ? 'route' : 'url',
+          dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+        },
+        ...traceparentData,
+      };
 
-        return maybePromiseResult;
-      })();
+      return trace(traceLoadContext, () => wrappingTarget.apply(thisArg, args), sendErrorToSentry);
     },
   });
 }
+
+function getTracePropagationData(event: ServerLoadEvent | LoadEvent): {
+  traceparentData?: TraceparentData;
+  dynamicSamplingContext?: Partial<DynamicSamplingContext>;
+} {
+  if (!isServerOnlyLoad(event)) {
+    return {};
+  }
+
+  const sentryTraceHeader = event.request.headers.get('sentry-trace');
+  const baggageHeader = event.request.headers.get('baggage');
+  const traceparentData = sentryTraceHeader ? extractTraceparentData(sentryTraceHeader) : undefined;
+  const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
+
+  return { traceparentData, dynamicSamplingContext };
+}
+
+/**
+ * Our server-side wrapLoadWithSentry can be used to wrap two different kinds of `load` functions:
+ *  - load functions from `+(page|layout).ts`: These can be called both on client and on server
+ *  - load functions from `+(page|layout).server.ts`: These are only called on the server
+ *
+ * In both cases, load events look differently. We can distinguish them by checking if the
+ * event has a `request` field (which only the server-exclusive load event has).
+ */
+function isServerOnlyLoad(event: ServerLoadEvent | LoadEvent): event is ServerLoadEvent {
+  return 'request' in event;
+}

packages/sveltekit/test/server/load.test.ts

@@ -54,6 +54,15 @@ const MOCK_LOAD_ARGS: any = {
     id: '/users/[id]',
   },
   url: new URL('http://localhost:3000/users/123'),
+};
+
+const MOCK_LOAD_NO_ROUTE_ARGS: any = {
+  params: { id: '123' },
+  url: new URL('http://localhost:3000/users/123'),
+};
+
+const MOCK_SERVER_ONLY_LOAD_ARGS: any = {
+  ...MOCK_LOAD_ARGS,
   request: {
     headers: {
       get: (key: string) => {
@@ -75,6 +84,32 @@ const MOCK_LOAD_ARGS: any = {
   },
 };
 
+const MOCK_SERVER_ONLY_NO_TRACE_LOAD_ARGS: any = {
+  ...MOCK_LOAD_ARGS,
+  request: {
+    headers: {
+      get: (_: string) => {
+        return null;
+      },
+    },
+  },
+};
+
+const MOCK_SERVER_ONLY_NO_BAGGAGE_LOAD_ARGS: any = {
+  ...MOCK_LOAD_ARGS,
+  request: {
+    headers: {
+      get: (key: string) => {
+        if (key === 'sentry-trace') {
+          return '1234567890abcdef1234567890abcdef-1234567890abcdef-1';
+        }
+
+        return null;
+      },
+    },
+  },
+};
+
 beforeAll(() => {
   addTracingExtensions();
 });
@@ -101,44 +136,6 @@ describe('wrapLoadWithSentry', () => {
     expect(mockCaptureException).toHaveBeenCalledTimes(1);
   });
 
-  // TODO: enable this once we figured out how tracing the load function doesn't result in creating a new transaction
-  it.skip('calls trace function', async () => {
-    async function load({ params }: Parameters<ServerLoad>[0]): Promise<ReturnType<ServerLoad>> {
-      return {
-        post: params.id,
-      };
-    }
-
-    const wrappedLoad = wrapLoadWithSentry(load);
-    await wrappedLoad(MOCK_LOAD_ARGS);
-
-    expect(mockTrace).toHaveBeenCalledTimes(1);
-    expect(mockTrace).toHaveBeenCalledWith(
-      {
-        op: 'function.sveltekit.load',
-        name: '/users/[id]',
-        parentSampled: true,
-        parentSpanId: '1234567890abcdef',
-        status: 'ok',
-        traceId: '1234567890abcdef1234567890abcdef',
-        metadata: {
-          dynamicSamplingContext: {
-            environment: 'production',
-            public_key: 'dogsarebadatkeepingsecrets',
-            release: '1.0.0',
-            sample_rate: '1',
-            trace_id: '1234567890abcdef1234567890abcdef',
-            transaction: 'dogpark',
-            user_segment: 'segmentA',
-          },
-          source: 'route',
-        },
-      },
-      expect.any(Function),
-      expect.any(Function),
-    );
-  });
-
   describe('with error() helper', () => {
     it.each([
       // [statusCode, timesCalled]
@@ -189,4 +186,125 @@ describe('wrapLoadWithSentry', () => {
       { handled: false, type: 'sveltekit', data: { function: 'load' } },
     );
   });
+
+  describe('calls trace', () => {
+    async function load({ params }: Parameters<ServerLoad>[0]): Promise<ReturnType<ServerLoad>> {
+      return {
+        post: params.id,
+      };
+    }
+
+    describe('for server-only load', () => {
+      it('attaches trace data if available', async () => {
+        const wrappedLoad = wrapLoadWithSentry(load);
+        await wrappedLoad(MOCK_SERVER_ONLY_LOAD_ARGS);
+
+        expect(mockTrace).toHaveBeenCalledTimes(1);
+        expect(mockTrace).toHaveBeenCalledWith(
+          {
+            op: 'function.sveltekit.load',
+            name: '/users/[id]',
+            parentSampled: true,
+            parentSpanId: '1234567890abcdef',
+            status: 'ok',
+            traceId: '1234567890abcdef1234567890abcdef',
+            metadata: {
+              dynamicSamplingContext: {
+                environment: 'production',
+                public_key: 'dogsarebadatkeepingsecrets',
+                release: '1.0.0',
+                sample_rate: '1',
+                trace_id: '1234567890abcdef1234567890abcdef',
+                transaction: 'dogpark',
+                user_segment: 'segmentA',
+              },
+              source: 'route',
+            },
+          },
+          expect.any(Function),
+          expect.any(Function),
+        );
+      });
+
+      it("doesn't attach trace data if it's not available", async () => {
+        const wrappedLoad = wrapLoadWithSentry(load);
+        await wrappedLoad(MOCK_SERVER_ONLY_NO_TRACE_LOAD_ARGS);
+
+        expect(mockTrace).toHaveBeenCalledTimes(1);
+        expect(mockTrace).toHaveBeenCalledWith(
+          {
+            op: 'function.sveltekit.load',
+            name: '/users/[id]',
+            status: 'ok',
+            metadata: {
+              source: 'route',
+            },
+          },
+          expect.any(Function),
+          expect.any(Function),
+        );
+      });
+
+      it("doesn't attach the DSC data if the baggage header not available", async () => {
+        const wrappedLoad = wrapLoadWithSentry(load);
+        await wrappedLoad(MOCK_SERVER_ONLY_NO_BAGGAGE_LOAD_ARGS);
+
+        expect(mockTrace).toHaveBeenCalledTimes(1);
+        expect(mockTrace).toHaveBeenCalledWith(
+          {
+            op: 'function.sveltekit.load',
+            name: '/users/[id]',
+            parentSampled: true,
+            parentSpanId: '1234567890abcdef',
+            status: 'ok',
+            traceId: '1234567890abcdef1234567890abcdef',
+            metadata: {
+              dynamicSamplingContext: {},
+              source: 'route',
+            },
+          },
+          expect.any(Function),
+          expect.any(Function),
+        );
+      });
+    });
+
+    it('for shared load', async () => {
+      const wrappedLoad = wrapLoadWithSentry(load);
+      await wrappedLoad(MOCK_LOAD_ARGS);
+
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+      expect(mockTrace).toHaveBeenCalledWith(
+        {
+          op: 'function.sveltekit.load',
+          name: '/users/[id]',
+          status: 'ok',
+          metadata: {
+            source: 'route',
+          },
+        },
+        expect.any(Function),
+        expect.any(Function),
+      );
+    });
+
+    it('falls back to the raw url if `event.route.id` is not available', async () => {
+      const wrappedLoad = wrapLoadWithSentry(load);
+      await wrappedLoad(MOCK_LOAD_NO_ROUTE_ARGS);
+
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+      expect(mockTrace).toHaveBeenCalledWith(
+        {
+          op: 'function.sveltekit.load',
+          name: '/users/123',
+          status: 'ok',
+          metadata: {
+            source: 'url',
+          },
+        },
+        expect.any(Function),
+        expect.any(Function),
+      );
+    });
+  });
 });