fix(replay): Ensure replay_id is removed from frozen DSC

getsentry · Oct 7, 2024 · ea7ce30 · ea7ce30
1 parent 3f0926e
commit ea7ce30
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 0 deletions.
diff --git a/packages/core/src/baseclient.ts b/packages/core/src/baseclient.ts
@@ -430,6 +430,9 @@ export abstract class BaseClient<O extends ClientOptions> implements Client<O> {
   /** @inheritdoc */
   public on(hook: 'createDsc', callback: (dsc: DynamicSamplingContext, rootSpan?: Span) => void): () => void;
 
+  /** @inheritdoc */
+  public on(hook: 'useFrozenDsc', callback: (dsc: Partial<DynamicSamplingContext>, rootSpan: Span) => void): () => void;
+
   /** @inheritdoc */
   public on(
     hook: 'beforeSendFeedback',
@@ -527,6 +530,9 @@ export abstract class BaseClient<O extends ClientOptions> implements Client<O> {
   /** @inheritdoc */
   public emit(hook: 'createDsc', dsc: DynamicSamplingContext, rootSpan?: Span): void;
 
+  /** @inheritdoc */
+  public emit(hook: 'useFrozenDsc', dsc: Partial<DynamicSamplingContext>, rootSpan: Span): void;
+
   /** @inheritdoc */
   public emit(hook: 'beforeSendFeedback', feedback: FeedbackEvent, options?: { includeReplay: boolean }): void;
 

diff --git a/packages/core/src/tracing/dynamicSamplingContext.ts b/packages/core/src/tracing/dynamicSamplingContext.ts
@@ -72,6 +72,7 @@ export function getDynamicSamplingContextFromSpan(span: Span): Readonly<Partial<
   // For core implementation, we freeze the DSC onto the span as a non-enumerable property
   const frozenDsc = (rootSpan as SpanWithMaybeDsc)[FROZEN_DSC_FIELD];
   if (frozenDsc) {
+    client.emit('useFrozenDsc', frozenDsc, rootSpan);
     return frozenDsc;
   }
 
@@ -83,6 +84,7 @@ export function getDynamicSamplingContextFromSpan(span: Span): Readonly<Partial<
   const dscOnTraceState = traceStateDsc && baggageHeaderToDynamicSamplingContext(traceStateDsc);
 
   if (dscOnTraceState) {
+    client.emit('useFrozenDsc', dscOnTraceState, rootSpan);
     return dscOnTraceState;
   }
 

diff --git a/packages/replay-internal/src/util/addGlobalListeners.ts b/packages/replay-internal/src/util/addGlobalListeners.ts
@@ -47,6 +47,26 @@ export function addGlobalListeners(replay: ReplayContainer): void {
       }
     });
 
+    // If a frozen DSC is used from a span, we want to make sure to remove the replay ID
+    // if the session has expired in the meanwhile
+    client.on('useFrozenDsc', (dsc: Partial<DynamicSamplingContext>) => {
+      if (!dsc.replay_id) {
+        return;
+      }
+
+      const replayId = replay.getSessionId();
+      if (
+        !replayId ||
+        !replay.isEnabled() ||
+        // We do not want to set the DSC when in buffer mode, as that means the replay has not been sent (yet)
+        replay.recordingMode !== 'session' ||
+        // This returns false if the session has expired in the meanwhile
+        !replay.checkAndHandleExpiredSession()
+      ) {
+        delete dsc.replay_id;
+      }
+    });
+
     client.on('spanStart', span => {
       replay.lastActiveSpan = span;
     });

diff --git a/packages/types/src/client.ts b/packages/types/src/client.ts
@@ -262,6 +262,12 @@ export interface Client<O extends ClientOptions = ClientOptions> {
    */
   on(hook: 'createDsc', callback: (dsc: DynamicSamplingContext, rootSpan?: Span) => void): () => void;
 
+  /**
+   * Register a callback to be used whenever a frozen DSC is used.
+   * @returns A function that, when executed, removes the registered callback.
+   */
+  on(hook: 'useFrozenDsc', callback: (dsc: Partial<DynamicSamplingContext>, rootSpan: Span) => void): () => void;
+
   /**
    * Register a callback when a Feedback event has been prepared.
    * This should be used to mutate the event. The options argument can hint
@@ -366,6 +372,12 @@ export interface Client<O extends ClientOptions = ClientOptions> {
    */
   emit(hook: 'createDsc', dsc: DynamicSamplingContext, rootSpan?: Span): void;
 
+  /**
+   * Fire a hook for when a frozen DSC (Dynamic Sampling Context) from a span is used.
+   * Expects the DSC as second argument, and the root span as third.
+   */
+  emit(hook: 'useFrozenDsc', dsc: Partial<DynamicSamplingContext>, rootSpan: Span): void;
+
   /**
    * Fire a hook event for after preparing a feedback event. Events to be given
    * a feedback event as the second argument, and an optional options object as