feat(core): Add beforeSendSpan hook (#11886) (#11918)

andreiborza · web-flow · commit f1d1a16952b8 · 2024-05-16T15:56:33.000+02:00
diff --git a/packages/core/src/baseclient.ts b/packages/core/src/baseclient.ts
@@ -24,6 +24,7 @@ import type {
   Span,
   SpanAttributes,
   SpanContextData,
+  SpanJSON,
   StartSpanOptions,
   TransactionEvent,
   Transport,
@@ -896,14 +897,27 @@ function processBeforeSend(
   event: Event,
   hint: EventHint,
 ): PromiseLike<Event | null> | Event | null {
-  const { beforeSend, beforeSendTransaction } = options;
+  const { beforeSend, beforeSendTransaction, beforeSendSpan } = options;
 
   if (isErrorEvent(event) && beforeSend) {
     return beforeSend(event, hint);
   }
 
-  if (isTransactionEvent(event) && beforeSendTransaction) {
-    return beforeSendTransaction(event, hint);
+  if (isTransactionEvent(event)) {
+    if (event.spans && beforeSendSpan) {
+      const processedSpans: SpanJSON[] = [];
+      for (const span of event.spans) {
+        const processedSpan = beforeSendSpan(span);
+        if (processedSpan) {
+          processedSpans.push(processedSpan);
+        }
+      }
+      event.spans = processedSpans;
+    }
+
+    if (beforeSendTransaction) {
+      return beforeSendTransaction(event, hint);
+    }
   }
 
   return event;
diff --git a/packages/core/src/envelope.ts b/packages/core/src/envelope.ts
@@ -1,4 +1,5 @@
 import type {
+  Client,
   DsnComponents,
   DynamicSamplingContext,
   Event,
@@ -11,6 +12,8 @@ import type {
   SessionEnvelope,
   SessionItem,
   SpanEnvelope,
+  SpanItem,
+  SpanJSON,
 } from '@sentry/types';
 import {
   createEnvelope,
@@ -94,8 +97,10 @@ export function createEventEnvelope(
 
 /**
  * Create envelope from Span item.
+ *
+ * Takes an optional client and runs spans through `beforeSendSpan` if available.
  */
-export function createSpanEnvelope(spans: SentrySpan[]): SpanEnvelope {
+export function createSpanEnvelope(spans: SentrySpan[], client?: Client): SpanEnvelope {
   function dscHasRequiredProps(dsc: Partial<DynamicSamplingContext>): dsc is DynamicSamplingContext {
     return !!dsc.trace_id && !!dsc.public_key;
   }
@@ -109,6 +114,19 @@ export function createSpanEnvelope(spans: SentrySpan[]): SpanEnvelope {
     sent_at: new Date().toISOString(),
     ...(dscHasRequiredProps(dsc) && { trace: dsc }),
   };
-  const items = spans.map(span => createSpanEnvelopeItem(spanToJSON(span)));
+
+  const beforeSendSpan = client && client.getOptions().beforeSendSpan;
+  const convertToSpanJSON = beforeSendSpan
+    ? (span: SentrySpan) => beforeSendSpan(spanToJSON(span) as SpanJSON)
+    : (span: SentrySpan) => spanToJSON(span);
+
+  const items: SpanItem[] = [];
+  for (const span of spans) {
+    const spanJson = convertToSpanJSON(span);
+    if (spanJson) {
+      items.push(createSpanEnvelopeItem(spanJson));
+    }
+  }
+
   return createEnvelope<SpanEnvelope>(headers, items);
 }
diff --git a/packages/core/src/tracing/sentrySpan.ts b/packages/core/src/tracing/sentrySpan.ts
@@ -97,12 +97,12 @@ export class SentrySpan implements Span {
 
     this._events = [];
 
+    this._isStandaloneSpan = spanContext.isStandalone;
+
     // If the span is already ended, ensure we finalize the span immediately
     if (this._endTime) {
       this._onSpanEnded();
     }
-
-    this._isStandaloneSpan = spanContext.isStandalone;
   }
 
   /** @inheritdoc */
@@ -259,7 +259,7 @@ export class SentrySpan implements Span {
 
     // if this is a standalone span, we send it immediately
     if (this._isStandaloneSpan) {
-      sendSpanEnvelope(createSpanEnvelope([this]));
+      sendSpanEnvelope(createSpanEnvelope([this], client));
       return;
     }
 
@@ -357,12 +357,24 @@ function isStandaloneSpan(span: Span): boolean {
   return span instanceof SentrySpan && span.isStandaloneSpan();
 }
 
+/**
+ * Sends a `SpanEnvelope`.
+ *
+ * Note: If the envelope's spans are dropped, e.g. via `beforeSendSpan`,
+ * the envelope will not be sent either.
+ */
 function sendSpanEnvelope(envelope: SpanEnvelope): void {
   const client = getClient();
   if (!client) {
     return;
   }
 
+  const spanItems = envelope[1];
+  if (!spanItems || spanItems.length === 0) {
+    client.recordDroppedEvent('before_send', 'span');
+    return;
+  }
+
   const transport = client.getTransport();
   if (transport) {
     transport.send(envelope).then(null, reason => {
diff --git a/packages/core/test/lib/base.test.ts b/packages/core/test/lib/base.test.ts
@@ -978,6 +978,38 @@ describe('BaseClient', () => {
       expect(TestClient.instance!.event!.transaction).toBe('/dogs/are/great');
     });
 
+    test('calls `beforeSendSpan` and uses original spans without any changes', () => {
+      expect.assertions(2);
+
+      const beforeSendSpan = jest.fn(span => span);
+      const options = getDefaultTestClientOptions({ dsn: PUBLIC_DSN, beforeSendSpan });
+      const client = new TestClient(options);
+
+      const transaction: Event = {
+        transaction: '/cats/are/great',
+        type: 'transaction',
+        spans: [
+          {
+            description: 'first span',
+            span_id: '9e15bf99fbe4bc80',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+          {
+            description: 'second span',
+            span_id: 'aa554c1f506b0783',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+        ],
+      };
+      client.captureEvent(transaction);
+
+      expect(beforeSendSpan).toHaveBeenCalledTimes(2);
+      const capturedEvent = TestClient.instance!.event!;
+      expect(capturedEvent.spans).toEqual(transaction.spans);
+    });
+
     test('calls `beforeSend` and uses the modified event', () => {
       expect.assertions(2);
 
@@ -1010,6 +1042,45 @@ describe('BaseClient', () => {
       expect(TestClient.instance!.event!.transaction).toBe('/adopt/dont/shop');
     });
 
+    test('calls `beforeSendSpan` and uses the modified spans', () => {
+      expect.assertions(3);
+
+      const beforeSendSpan = jest.fn(span => {
+        span.data = { version: 'bravo' };
+        return span;
+      });
+
+      const options = getDefaultTestClientOptions({ dsn: PUBLIC_DSN, beforeSendSpan });
+      const client = new TestClient(options);
+      const transaction: Event = {
+        transaction: '/cats/are/great',
+        type: 'transaction',
+        spans: [
+          {
+            description: 'first span',
+            span_id: '9e15bf99fbe4bc80',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+          {
+            description: 'second span',
+            span_id: 'aa554c1f506b0783',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+        ],
+      };
+
+      client.captureEvent(transaction);
+
+      expect(beforeSendSpan).toHaveBeenCalledTimes(2);
+      const capturedEvent = TestClient.instance!.event!;
+      for (const [idx, span] of capturedEvent.spans!.entries()) {
+        const originalSpan = transaction.spans![idx];
+        expect(span).toEqual({ ...originalSpan, data: { version: 'bravo' } });
+      }
+    });
+
     test('calls `beforeSend` and discards the event', () => {
       expect.assertions(4);
 
@@ -1048,6 +1119,38 @@ describe('BaseClient', () => {
       expect(loggerWarnSpy).toBeCalledWith('before send for type `transaction` returned `null`, will not send event.');
     });
 
+    test('calls `beforeSendSpan` and discards the span', () => {
+      expect.assertions(2);
+
+      const beforeSendSpan = jest.fn(() => null);
+      const options = getDefaultTestClientOptions({ dsn: PUBLIC_DSN, beforeSendSpan });
+      const client = new TestClient(options);
+
+      const transaction: Event = {
+        transaction: '/cats/are/great',
+        type: 'transaction',
+        spans: [
+          {
+            description: 'first span',
+            span_id: '9e15bf99fbe4bc80',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+          {
+            description: 'second span',
+            span_id: 'aa554c1f506b0783',
+            start_timestamp: 1591603196.637835,
+            trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+          },
+        ],
+      };
+      client.captureEvent(transaction);
+
+      expect(beforeSendSpan).toHaveBeenCalledTimes(2);
+      const capturedEvent = TestClient.instance!.event!;
+      expect(capturedEvent.spans).toHaveLength(0);
+    });
+
     test('calls `beforeSend` and logs info about invalid return value', () => {
       const invalidValues = [undefined, false, true, [], 1];
       expect.assertions(invalidValues.length * 3);
diff --git a/packages/core/test/lib/envelope.test.ts b/packages/core/test/lib/envelope.test.ts
@@ -157,4 +157,71 @@ describe('createSpanEnvelope', () => {
       sent_at: expect.any(String),
     });
   });
+
+  it('calls `beforeSendSpan` and uses original span without any changes', () => {
+    const beforeSendSpan = jest.fn(span => span);
+    const options = getDefaultTestClientOptions({ dsn: 'https://domain/123', beforeSendSpan });
+    const client = new TestClient(options);
+
+    const span = new SentrySpan({
+      name: 'test',
+      isStandalone: true,
+      startTimestamp: 1,
+      endTimestamp: 2,
+    });
+
+    const spanEnvelope = createSpanEnvelope([span], client);
+
+    expect(beforeSendSpan).toHaveBeenCalled();
+
+    const spanItem = spanEnvelope[1][0][1];
+    expect(spanItem).toEqual({
+      data: {
+        'sentry.origin': 'manual',
+      },
+      description: 'test',
+      is_segment: true,
+      origin: 'manual',
+      span_id: expect.stringMatching(/^[0-9a-f]{16}$/),
+      segment_id: spanItem.segment_id,
+      start_timestamp: 1,
+      timestamp: 2,
+      trace_id: expect.stringMatching(/^[0-9a-f]{32}$/),
+    });
+  });
+
+  it('calls `beforeSendSpan` and uses the modified span', () => {
+    const beforeSendSpan = jest.fn(span => {
+      span.description = `mutated description: ${span.description}`;
+      return span;
+    });
+    const options = getDefaultTestClientOptions({ dsn: 'https://domain/123', beforeSendSpan });
+    const client = new TestClient(options);
+
+    const span = new SentrySpan({
+      name: 'test',
+      isStandalone: true,
+      startTimestamp: 1,
+      endTimestamp: 2,
+    });
+
+    const spanEnvelope = createSpanEnvelope([span], client);
+
+    expect(beforeSendSpan).toHaveBeenCalled();
+
+    const spanItem = spanEnvelope[1][0][1];
+    expect(spanItem).toEqual({
+      data: {
+        'sentry.origin': 'manual',
+      },
+      description: 'mutated description: test',
+      is_segment: true,
+      origin: 'manual',
+      span_id: expect.stringMatching(/^[0-9a-f]{16}$/),
+      segment_id: spanItem.segment_id,
+      start_timestamp: 1,
+      timestamp: 2,
+      trace_id: expect.stringMatching(/^[0-9a-f]{32}$/),
+    });
+  });
 });
diff --git a/packages/core/test/lib/tracing/sentrySpan.test.ts b/packages/core/test/lib/tracing/sentrySpan.test.ts
@@ -1,7 +1,9 @@
 import { timestampInSeconds } from '@sentry/utils';
+import { setCurrentClient } from '../../../src';
 import { SentrySpan } from '../../../src/tracing/sentrySpan';
 import { SPAN_STATUS_ERROR } from '../../../src/tracing/spanstatus';
 import { TRACE_FLAG_NONE, TRACE_FLAG_SAMPLED, spanToJSON } from '../../../src/utils/spanUtils';
+import { TestClient, getDefaultTestClientOptions } from '../../mocks/client';
 
 describe('SentrySpan', () => {
   describe('name', () => {
@@ -88,6 +90,57 @@ describe('SentrySpan', () => {
       span.end();
       expect(spanToJSON(span).timestamp).toBeGreaterThan(1);
     });
+
+    test('sends the span if `beforeSendSpan` does not modify the span ', () => {
+      const beforeSendSpan = jest.fn(span => span);
+      const client = new TestClient(
+        getDefaultTestClientOptions({
+          dsn: 'https://username@domain/123',
+          enableSend: true,
+          beforeSendSpan,
+        }),
+      );
+      setCurrentClient(client);
+
+      // @ts-expect-error Accessing private transport API
+      const mockSend = jest.spyOn(client._transport, 'send');
+      const span = new SentrySpan({
+        name: 'test',
+        isStandalone: true,
+        startTimestamp: 1,
+        endTimestamp: 2,
+        sampled: true,
+      });
+      span.end();
+      expect(mockSend).toHaveBeenCalled();
+    });
+
+    test('does not send the span if `beforeSendSpan` drops the span', () => {
+      const beforeSendSpan = jest.fn(() => null);
+      const client = new TestClient(
+        getDefaultTestClientOptions({
+          dsn: 'https://username@domain/123',
+          enableSend: true,
+          beforeSendSpan,
+        }),
+      );
+      setCurrentClient(client);
+
+      const recordDroppedEventSpy = jest.spyOn(client, 'recordDroppedEvent');
+      // @ts-expect-error Accessing private transport API
+      const mockSend = jest.spyOn(client._transport, 'send');
+      const span = new SentrySpan({
+        name: 'test',
+        isStandalone: true,
+        startTimestamp: 1,
+        endTimestamp: 2,
+        sampled: true,
+      });
+      span.end();
+
+      expect(mockSend).not.toHaveBeenCalled();
+      expect(recordDroppedEventSpy).toHaveBeenCalledWith('before_send', 'span');
+    });
   });
 
   describe('end', () => {
diff --git a/packages/types/src/options.ts b/packages/types/src/options.ts
