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

Lms24 · Lms24 · commit 85b56dcc5090 · 2023-03-23T12:07:57.000+01:00
diff --git a/packages/sveltekit/src/server/load.ts b/packages/sveltekit/src/server/load.ts
@@ -1,7 +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 type { TransactionContext } from '@sentry/types';
+import {
+  addExceptionMechanism,
+  baggageHeaderToDynamicSamplingContext,
+  extractTraceparentData,
+  objectify,
+} from '@sentry/utils';
+import type { HttpError, Load, LoadEvent, ServerLoad, ServerLoadEvent } from '@sveltejs/kit';
 import * as domain from 'domain';
 
 function isHttpError(err: unknown): err is HttpError {
@@ -43,27 +50,55 @@ function sendErrorToSentry(e: unknown): unknown {
  *
  * @param origLoad SvelteKit user defined load function
  */
-export function wrapLoadWithSentry<T extends ServerLoad | Load>(origLoad: T): T {
+export function wrapLoadWithSentry<T extends Load | ServerLoad>(origLoad: T): T {
   return new Proxy(origLoad, {
-    apply: (wrappingTarget, thisArg, args: Parameters<ServerLoad>) => {
+    apply: (wrappingTarget, thisArg, args: Parameters<ServerLoad | Load>) => {
       return domain.create().bind(() => {
-        let maybePromiseResult: ReturnType<T>;
+        const [event] = args;
+        const routeId = event.route && event.route.id;
 
-        try {
-          maybePromiseResult = wrappingTarget.apply(thisArg, args);
-        } catch (e) {
-          sendErrorToSentry(e);
-          throw e;
-        }
+        const traceSharedLoadContext: TransactionContext = {
+          op: 'function.sveltekit.load',
+          name: routeId ? routeId : event.url.pathname,
+          status: 'ok',
+          metadata: {
+            source: routeId ? 'route' : 'url',
+          },
+        };
+
+        let finalTraceLoadContext = { ...traceSharedLoadContext };
+
+        if (isServerOnlyLoad(event)) {
+          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);
 
-        if (isThenable(maybePromiseResult)) {
-          Promise.resolve(maybePromiseResult).then(null, e => {
-            sendErrorToSentry(e);
-          });
+          const traceSeverOnlyLoadContext = {
+            ...traceparentData,
+            metadata: {
+              ...traceSharedLoadContext.metadata,
+              dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
+            },
+          };
+
+          finalTraceLoadContext = { ...traceSharedLoadContext, ...traceSeverOnlyLoadContext };
         }
 
-        return maybePromiseResult;
+        return trace(finalTraceLoadContext, () => wrappingTarget.apply(thisArg, args), sendErrorToSentry);
       })();
     },
   });
 }
+
+/**
+ * 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;
+}
diff --git a/packages/sveltekit/test/server/load.test.ts b/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,42 +136,125 @@ 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 () => {
+  describe('calls trace', () => {
     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',
+    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',
           },
-          source: 'route',
         },
-      },
-      expect.any(Function),
-      expect.any(Function),
-    );
+        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),
+      );
+    });
   });
 
   describe('with error() helper', () => {
