feat(node): Add trace context to checkin

AbhiPrasad · AbhiPrasad · commit 5d10d1fb3f58 · 2023-07-11T14:20:35.000-04:00
diff --git a/packages/core/src/checkin.ts b/packages/core/src/checkin.ts
@@ -1,26 +1,42 @@
-import type { CheckInEvelope, CheckInItem, DsnComponents, SdkMetadata, SerializedCheckIn } from '@sentry/types';
-import { createEnvelope, dsnToString } from '@sentry/utils';
+import type {
+  CheckInEvelope,
+  CheckInItem,
+  DsnComponents,
+  DynamicSamplingContext,
+  SdkMetadata,
+  SerializedCheckIn,
+} from '@sentry/types';
+import { createEnvelope, dropUndefinedKeys, dsnToString } from '@sentry/utils';
 
 /**
  * Create envelope from check in item.
  */
 export function createCheckInEnvelope(
   checkIn: SerializedCheckIn,
+  dynamicSamplingContext?: Partial<DynamicSamplingContext>,
   metadata?: SdkMetadata,
   tunnel?: string,
   dsn?: DsnComponents,
 ): CheckInEvelope {
   const headers: CheckInEvelope[0] = {
     sent_at: new Date().toISOString(),
-    ...(metadata &&
-      metadata.sdk && {
-        sdk: {
-          name: metadata.sdk.name,
-          version: metadata.sdk.version,
-        },
-      }),
-    ...(!!tunnel && !!dsn && { dsn: dsnToString(dsn) }),
   };
+
+  if (metadata && metadata.sdk) {
+    headers.sdk = {
+      name: metadata.sdk.name,
+      version: metadata.sdk.version,
+    };
+  }
+
+  if (!!tunnel && !!dsn) {
+    headers.dsn = dsnToString(dsn);
+  }
+
+  if (dynamicSamplingContext) {
+    headers.trace = dropUndefinedKeys(dynamicSamplingContext) as DynamicSamplingContext;
+  }
+
   const item = createCheckInEnvelopeItem(checkIn);
   return createEnvelope<CheckInEvelope>(headers, [item]);
 }
diff --git a/packages/core/src/exports.ts b/packages/core/src/exports.ts
@@ -196,13 +196,15 @@ export function startTransaction(
  * to create a monitor automatically when sending a check in.
  */
 export function captureCheckIn(checkIn: CheckIn, upsertMonitorConfig?: MonitorConfig): string {
-  const client = getCurrentHub().getClient();
+  const hub = getCurrentHub();
+  const scope = hub.getScope();
+  const client = hub.getClient();
   if (!client) {
     __DEBUG_BUILD__ && logger.warn('Cannot capture check-in. No client defined.');
   } else if (!client.captureCheckIn) {
     __DEBUG_BUILD__ && logger.warn('Cannot capture check-in. Client does not support sending check-ins.');
   } else {
-    return client.captureCheckIn(checkIn, upsertMonitorConfig);
+    return client.captureCheckIn(checkIn, upsertMonitorConfig, scope);
   }
 
   return uuid4();
diff --git a/packages/core/test/lib/checkin.test.ts b/packages/core/test/lib/checkin.test.ts
@@ -10,6 +10,10 @@ describe('createCheckInEnvelope', () => {
         monitor_slug: 'b7645b8e-b47d-4398-be9a-d16b0dac31cb',
         status: 'in_progress',
       },
+      {
+        trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+        public_key: 'testPublicKey',
+      },
       {
         sdk: {
           name: 'testSdkName',
@@ -30,6 +34,10 @@ describe('createCheckInEnvelope', () => {
         name: 'testSdkName',
         version: 'testSdkVersion',
       },
+      trace: {
+        trace_id: '86f39e84263a4de99c326acab3bfe3bd',
+        public_key: 'testPublicKey',
+      },
       sent_at: expect.any(String),
     });
   });
diff --git a/packages/nextjs/src/edge/edgeclient.ts b/packages/nextjs/src/edge/edgeclient.ts
@@ -1,14 +1,22 @@
 import type { Scope } from '@sentry/core';
-import { addTracingExtensions, BaseClient, createCheckInEnvelope, SDK_VERSION } from '@sentry/core';
+import {
+  addTracingExtensions,
+  BaseClient,
+  createCheckInEnvelope,
+  getDynamicSamplingContextFromClient,
+  SDK_VERSION,
+} from '@sentry/core';
 import type {
   CheckIn,
   ClientOptions,
+  DynamicSamplingContext,
   Event,
   EventHint,
   MonitorConfig,
   SerializedCheckIn,
   Severity,
   SeverityLevel,
+  TraceContext,
 } from '@sentry/types';
 import { logger, uuid4 } from '@sentry/utils';
 
@@ -72,7 +80,7 @@ export class EdgeClient extends BaseClient<EdgeClientOptions> {
    * @param upsertMonitorConfig An optional object that describes a monitor config. Use this if you want
    * to create a monitor automatically when sending a check in.
    */
-  public captureCheckIn(checkIn: CheckIn, monitorConfig?: MonitorConfig): string {
+  public captureCheckIn(checkIn: CheckIn, monitorConfig?: MonitorConfig, scope?: Scope): string {
     const id = checkIn.status !== 'in_progress' && checkIn.checkInId ? checkIn.checkInId : uuid4();
     if (!this._isEnabled()) {
       __DEBUG_BUILD__ && logger.warn('SDK not enabled, will not capture checkin.');
@@ -103,7 +111,20 @@ export class EdgeClient extends BaseClient<EdgeClientOptions> {
       };
     }
 
-    const envelope = createCheckInEnvelope(serializedCheckIn, this.getSdkMetadata(), tunnel, this.getDsn());
+    const [dynamicSamplingContext, traceContext] = this._getTraceInfoFromScope(scope);
+    if (traceContext) {
+      serializedCheckIn.contexts = {
+        trace: traceContext,
+      };
+    }
+
+    const envelope = createCheckInEnvelope(
+      serializedCheckIn,
+      dynamicSamplingContext,
+      this.getSdkMetadata(),
+      tunnel,
+      this.getDsn(),
+    );
 
     __DEBUG_BUILD__ && logger.info('Sending checkin:', checkIn.monitorSlug, checkIn.status);
     void this._sendEnvelope(envelope);
@@ -124,4 +145,30 @@ export class EdgeClient extends BaseClient<EdgeClientOptions> {
     event.server_name = event.server_name || process.env.SENTRY_NAME;
     return super._prepareEvent(event, hint, scope);
   }
+
+  /** Extract trace information from scope */
+  private _getTraceInfoFromScope(
+    scope: Scope | undefined,
+  ): [dynamicSamplingContext: Partial<DynamicSamplingContext> | undefined, traceContext: TraceContext | undefined] {
+    if (!scope) {
+      return [undefined, undefined];
+    }
+
+    const span = scope.getSpan();
+    if (span) {
+      return [span?.transaction?.getDynamicSamplingContext(), span?.getTraceContext()];
+    }
+
+    const { traceId, spanId, parentSpanId, dsc } = scope.getPropagationContext();
+    const traceContext: TraceContext = {
+      trace_id: traceId,
+      span_id: spanId,
+      parent_span_id: parentSpanId,
+    };
+    if (dsc) {
+      return [dsc, traceContext];
+    }
+
+    return [getDynamicSamplingContextFromClient(traceId, this, scope), traceContext];
+  }
 }
diff --git a/packages/node/src/client.ts b/packages/node/src/client.ts
@@ -1,13 +1,22 @@
 import type { Scope } from '@sentry/core';
-import { addTracingExtensions, BaseClient, createCheckInEnvelope, SDK_VERSION, SessionFlusher } from '@sentry/core';
+import {
+  addTracingExtensions,
+  BaseClient,
+  createCheckInEnvelope,
+  getDynamicSamplingContextFromClient,
+  SDK_VERSION,
+  SessionFlusher,
+} from '@sentry/core';
 import type {
   CheckIn,
+  DynamicSamplingContext,
   Event,
   EventHint,
   MonitorConfig,
   SerializedCheckIn,
   Severity,
   SeverityLevel,
+  TraceContext,
 } from '@sentry/types';
 import { logger, resolvedSyncPromise, uuid4 } from '@sentry/utils';
 import * as os from 'os';
@@ -154,7 +163,7 @@ export class NodeClient extends BaseClient<NodeClientOptions> {
    * to create a monitor automatically when sending a check in.
    * @returns A string representing the id of the check in.
    */
-  public captureCheckIn(checkIn: CheckIn, monitorConfig?: MonitorConfig): string {
+  public captureCheckIn(checkIn: CheckIn, monitorConfig?: MonitorConfig, scope?: Scope): string {
     const id = checkIn.status !== 'in_progress' && checkIn.checkInId ? checkIn.checkInId : uuid4();
     if (!this._isEnabled()) {
       __DEBUG_BUILD__ && logger.warn('SDK not enabled, will not capture checkin.');
@@ -185,7 +194,20 @@ export class NodeClient extends BaseClient<NodeClientOptions> {
       };
     }
 
-    const envelope = createCheckInEnvelope(serializedCheckIn, this.getSdkMetadata(), tunnel, this.getDsn());
+    const [dynamicSamplingContext, traceContext] = this._getTraceInfoFromScope(scope);
+    if (traceContext) {
+      serializedCheckIn.contexts = {
+        trace: traceContext,
+      };
+    }
+
+    const envelope = createCheckInEnvelope(
+      serializedCheckIn,
+      dynamicSamplingContext,
+      this.getSdkMetadata(),
+      tunnel,
+      this.getDsn(),
+    );
 
     __DEBUG_BUILD__ && logger.info('Sending checkin:', checkIn.monitorSlug, checkIn.status);
     void this._sendEnvelope(envelope);
@@ -220,4 +242,30 @@ export class NodeClient extends BaseClient<NodeClientOptions> {
       this._sessionFlusher.incrementSessionStatusCount();
     }
   }
+
+  /** Extract trace information from scope */
+  private _getTraceInfoFromScope(
+    scope: Scope | undefined,
+  ): [dynamicSamplingContext: Partial<DynamicSamplingContext> | undefined, traceContext: TraceContext | undefined] {
+    if (!scope) {
+      return [undefined, undefined];
+    }
+
+    const span = scope.getSpan();
+    if (span) {
+      return [span?.transaction?.getDynamicSamplingContext(), span?.getTraceContext()];
+    }
+
+    const { traceId, spanId, parentSpanId, dsc } = scope.getPropagationContext();
+    const traceContext: TraceContext = {
+      trace_id: traceId,
+      span_id: spanId,
+      parent_span_id: parentSpanId,
+    };
+    if (dsc) {
+      return [dsc, traceContext];
+    }
+
+    return [getDynamicSamplingContextFromClient(traceId, this, scope), traceContext];
+  }
 }
diff --git a/packages/types/src/checkin.ts b/packages/types/src/checkin.ts
@@ -1,3 +1,5 @@
+import type { TraceContext } from './context';
+
 interface CrontabSchedule {
   type: 'crontab';
   // The crontab schedule string, e.g. 0 * * * *.
@@ -36,6 +38,9 @@ export interface SerializedCheckIn {
     // See: https://en.wikipedia.org/wiki/List_of_tz_database_time_zones
     timezone?: string;
   };
+  contexts?: {
+    trace?: TraceContext;
+  };
 }
 
 interface InProgressCheckIn {
diff --git a/packages/types/src/client.ts b/packages/types/src/client.ts
@@ -74,9 +74,10 @@ export interface Client<O extends ClientOptions = ClientOptions> {
    * @param checkIn An object that describes a check in.
    * @param upsertMonitorConfig An optional object that describes a monitor config. Use this if you want
    * to create a monitor automatically when sending a check in.
+   * @param scope An optional scope containing event metadata.
    * @returns A string representing the id of the check in.
    */
-  captureCheckIn?(checkIn: CheckIn, monitorConfig?: MonitorConfig): string;
+  captureCheckIn?(checkIn: CheckIn, monitorConfig?: MonitorConfig, scope?: Scope): string;
 
   /** Returns the current Dsn. */
   getDsn(): DsnComponents | undefined;
diff --git a/packages/types/src/envelope.ts b/packages/types/src/envelope.ts
@@ -86,7 +86,7 @@ type ReplayRecordingItem = BaseEnvelopeItem<ReplayRecordingItemHeaders, ReplayRe
 
 export type EventEnvelopeHeaders = { event_id: string; sent_at: string; trace?: DynamicSamplingContext };
 type SessionEnvelopeHeaders = { sent_at: string };
-type CheckInEnvelopeHeaders = BaseEnvelopeHeaders;
+type CheckInEnvelopeHeaders = { trace?: DynamicSamplingContext };
 type ClientReportEnvelopeHeaders = BaseEnvelopeHeaders;
 type ReplayEnvelopeHeaders = BaseEnvelopeHeaders;
 
