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

Lms24 · web-flow · commit 0a4a07266bf5 · 2023-03-23T17:01:45.000+01:00
Previously, our server `wrapLoadWithSentry` wrapper assumed that `load` in `+page.ts` and `+page.server.ts` work identically. Unfortunately, the load `event` arg has different contents for both, namely that the server-only load event contains a `request` object while the shared (client and server) load event doesn't. This patch fixes that by distinguishing between the two load types within the wrapper. Users can still use the same wrapper for both cases.

Furthermore, this patch removes the usage of domains in the load wrapper, as on the server-side, as we're already in a domain when `handleSentry` is used. Creating another domain here, caused the creation of a new transaction instead of adding the span to the transaction that was created by `handleSentry`.
diff --git a/packages/sveltekit/src/server/load.ts b/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;
+}
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,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),
+      );
+    });
+  });
 });
