fix(serverless): Handle incoming "sentry-trace" header (#3261)

* Applying sentry-trace in AWSLambdaIntegration transaction

* Propagate trace in AWSLambda functions

* Propagate trace handlers in GCP http functions

* Linting fix

* Linting fixes

* Removed redundant type

Author: ahmedetefy

packages/serverless/package.json

@@ -18,6 +18,7 @@
   "dependencies": {
     "@sentry/minimal": "6.1.0",
     "@sentry/node": "6.1.0",
+    "@sentry/tracing": "6.1.0",
     "@sentry/types": "6.1.0",
     "@sentry/utils": "6.1.0",
     "@types/aws-lambda": "^8.10.62",

packages/serverless/src/awslambda.ts

@@ -9,8 +9,9 @@ import {
   withScope,
 } from '@sentry/node';
 import * as Sentry from '@sentry/node';
+import { extractTraceparentData } from '@sentry/tracing';
 import { Integration } from '@sentry/types';
-import { logger } from '@sentry/utils';
+import { isString, logger } from '@sentry/utils';
 // NOTE: I have no idea how to fix this right now, and don't want to waste more time, as it builds just fine — Kamil
 // eslint-disable-next-line import/no-unresolved
 import { Context, Handler } from 'aws-lambda';
@@ -242,9 +243,16 @@ export function wrapHandler<TEvent, TResult>(
       }, timeoutWarningDelay);
     }
 
+    // Applying `sentry-trace` to context
+    let traceparentData;
+    const eventWithHeaders = event as { headers?: { [key: string]: string } };
+    if (eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])) {
+      traceparentData = extractTraceparentData(eventWithHeaders.headers['sentry-trace'] as string);
+    }
     const transaction = startTransaction({
       name: context.functionName,
       op: 'awslambda.handler',
+      ...traceparentData,
     });
 
     const hub = getCurrentHub();

packages/serverless/src/gcpfunction/http.ts

@@ -1,5 +1,6 @@
 import { captureException, flush, getCurrentHub, Handlers, startTransaction } from '@sentry/node';
-import { logger, stripUrlQueryAndFragment } from '@sentry/utils';
+import { extractTraceparentData } from '@sentry/tracing';
+import { isString, logger, stripUrlQueryAndFragment } from '@sentry/utils';
 
 import { domainify, getActiveDomain, proxyFunction } from './../utils';
 import { HttpFunction, WrapperOptions } from './general';
@@ -48,9 +49,16 @@ function _wrapHttpFunction(fn: HttpFunction, wrapOptions: Partial<HttpFunctionWr
     const reqMethod = (req.method || '').toUpperCase();
     const reqUrl = stripUrlQueryAndFragment(req.originalUrl || req.url || '');
 
+    // Applying `sentry-trace` to context
+    let traceparentData;
+    const reqWithHeaders = req as { headers?: { [key: string]: string } };
+    if (reqWithHeaders.headers && isString(reqWithHeaders.headers['sentry-trace'])) {
+      traceparentData = extractTraceparentData(reqWithHeaders.headers['sentry-trace'] as string);
+    }
     const transaction = startTransaction({
       name: `${reqMethod} ${reqUrl}`,
       op: 'gcp.function.http',
+      ...traceparentData,
     });
 
     // getCurrentHub() is expected to use current active domain as a carrier

packages/serverless/test/awslambda.test.ts

@@ -1,4 +1,3 @@
-import { Event } from '@sentry/types';
 // NOTE: I have no idea how to fix this right now, and don't want to waste more time, as it builds just fine — Kamil
 // eslint-disable-next-line import/no-unresolved
 import { Callback, Handler } from 'aws-lambda';
@@ -16,9 +15,7 @@ const { wrapHandler } = Sentry.AWSLambda;
 
 // Default `timeoutWarningLimit` is 500ms so leaving some space for it to trigger when necessary
 const DEFAULT_EXECUTION_TIME = 100;
-const fakeEvent = {
-  fortySix: 'o_O',
-};
+let fakeEvent: { [key: string]: unknown };
 const fakeContext = {
   callbackWaitsForEmptyEventLoop: false,
   functionName: 'functionName',
@@ -72,6 +69,12 @@ function expectScopeSettings() {
 }
 
 describe('AWSLambda', () => {
+  beforeEach(() => {
+    fakeEvent = {
+      fortySix: 'o_O',
+    };
+  });
+
   afterEach(() => {
     // @ts-ignore see "Why @ts-ignore" note
     Sentry.resetMocks();
@@ -228,9 +231,16 @@ describe('AWSLambda', () => {
       const wrappedHandler = wrapHandler(handler);
 
       try {
+        fakeEvent.headers = { 'sentry-trace': '12312012123120121231201212312012-1121201211212012-0' };
         await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
       } catch (e) {
-        expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
+        expect(Sentry.startTransaction).toBeCalledWith({
+          name: 'functionName',
+          op: 'awslambda.handler',
+          traceId: '12312012123120121231201212312012',
+          parentSpanId: '1121201211212012',
+          parentSampled: false,
+        });
         expectScopeSettings();
         expect(Sentry.captureException).toBeCalledWith(e);
         // @ts-ignore see "Why @ts-ignore" note

packages/serverless/test/gcpfunction.test.ts

@@ -26,13 +26,17 @@ describe('GCPFunction', () => {
     Sentry.resetMocks();
   });
 
-  async function handleHttp(fn: HttpFunction): Promise<void> {
+  async function handleHttp(fn: HttpFunction, trace_headers: { [key: string]: string } | null = null): Promise<void> {
+    let headers: { [key: string]: string } = { host: 'hostname', 'content-type': 'application/json' };
+    if (trace_headers) {
+      headers = { ...headers, ...trace_headers };
+    }
     return new Promise((resolve, _reject) => {
       const d = domain.create();
       const req = {
         method: 'POST',
         url: '/path?q=query',
-        headers: { host: 'hostname', 'content-type': 'application/json' },
+        headers: headers,
         body: { foo: 'bar' },
       } as Request;
       const res = { end: resolve } as Response;
@@ -124,8 +128,19 @@ describe('GCPFunction', () => {
         throw error;
       };
       const wrappedHandler = wrapHttpFunction(handler);
-      await handleHttp(wrappedHandler);
-      expect(Sentry.startTransaction).toBeCalledWith({ name: 'POST /path', op: 'gcp.function.http' });
+
+      const trace_headers: { [key: string]: string } = {
+        'sentry-trace': '12312012123120121231201212312012-1121201211212012-0',
+      };
+
+      await handleHttp(wrappedHandler, trace_headers);
+      expect(Sentry.startTransaction).toBeCalledWith({
+        name: 'POST /path',
+        op: 'gcp.function.http',
+        traceId: '12312012123120121231201212312012',
+        parentSpanId: '1121201211212012',
+        parentSampled: false,
+      });
       // @ts-ignore see "Why @ts-ignore" note
       expect(Sentry.fakeScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
       expect(Sentry.captureException).toBeCalledWith(error);