feat(opentelemetry): Expose sampling helper (#12674)

When users want to use a custom sampler, they can use these new helpers
to still have sentry working nicely with whatever they decide to do in
there.

For e.g. trace propagation etc. to work correctly with Sentry, we need
to attach some things to trace state etc. These helpers encapsulate this
for the user, while still allowing them to decide however they want if
the span should be sampled or not.

This was brought up here:
#12191 (reply in thread)

Author: mydea

packages/opentelemetry/src/index.ts

@@ -36,7 +36,10 @@ export { setOpenTelemetryContextAsyncContextStrategy } from './asyncContextStrat
 export { wrapContextManagerClass } from './contextManager';
 export { SentryPropagator } from './propagator';
 export { SentrySpanProcessor } from './spanProcessor';
-export { SentrySampler } from './sampler';
+export {
+  SentrySampler,
+  wrapSamplingDecision,
+} from './sampler';
 
 export { openTelemetrySetupCheck } from './utils/setupCheck';
 

packages/opentelemetry/src/sampler.ts

@@ -1,4 +1,4 @@
-import type { Attributes, Context, Span } from '@opentelemetry/api';
+import type { Attributes, Context, Span, TraceState as TraceStateInterface } from '@opentelemetry/api';
 import { SpanKind } from '@opentelemetry/api';
 import { isSpanContextValid, trace } from '@opentelemetry/api';
 import { TraceState } from '@opentelemetry/core';
@@ -40,16 +40,8 @@ export class SentrySampler implements Sampler {
     const parentSpan = trace.getSpan(context);
     const parentContext = parentSpan?.spanContext();
 
-    let traceState = parentContext?.traceState || new TraceState();
-
-    // We always keep the URL on the trace state, so we can access it in the propagator
-    const url = spanAttributes[SEMATTRS_HTTP_URL];
-    if (url && typeof url === 'string') {
-      traceState = traceState.set(SENTRY_TRACE_STATE_URL, url);
-    }
-
     if (!hasTracingEnabled(options)) {
-      return { decision: SamplingDecision.NOT_RECORD, traceState };
+      return wrapSamplingDecision({ decision: undefined, context, spanAttributes });
     }
 
     // If we have a http.client span that has no local parent, we never want to sample it
@@ -59,7 +51,7 @@ export class SentrySampler implements Sampler {
       spanAttributes[SEMATTRS_HTTP_METHOD] &&
       (!parentSpan || parentContext?.isRemote)
     ) {
-      return { decision: SamplingDecision.NOT_RECORD, traceState };
+      return wrapSamplingDecision({ decision: undefined, context, spanAttributes });
     }
 
     const parentSampled = parentSpan ? getParentSampled(parentSpan, traceId, spanName) : undefined;
@@ -76,7 +68,7 @@ export class SentrySampler implements Sampler {
       mutableSamplingDecision,
     );
     if (!mutableSamplingDecision.decision) {
-      return { decision: SamplingDecision.NOT_RECORD, traceState: traceState };
+      return wrapSamplingDecision({ decision: undefined, context, spanAttributes });
     }
 
     const [sampled, sampleRate] = sampleSpan(options, {
@@ -96,25 +88,22 @@ export class SentrySampler implements Sampler {
     const method = `${spanAttributes[SEMATTRS_HTTP_METHOD]}`.toUpperCase();
     if (method === 'OPTIONS' || method === 'HEAD') {
       DEBUG_BUILD && logger.log(`[Tracing] Not sampling span because HTTP method is '${method}' for ${spanName}`);
+
       return {
-        decision: SamplingDecision.NOT_RECORD,
+        ...wrapSamplingDecision({ decision: SamplingDecision.NOT_RECORD, context, spanAttributes }),
         attributes,
-        traceState: traceState.set(SENTRY_TRACE_STATE_SAMPLED_NOT_RECORDING, '1'),
       };
     }
 
     if (!sampled) {
       return {
-        decision: SamplingDecision.NOT_RECORD,
+        ...wrapSamplingDecision({ decision: SamplingDecision.NOT_RECORD, context, spanAttributes }),
         attributes,
-        traceState: traceState.set(SENTRY_TRACE_STATE_SAMPLED_NOT_RECORDING, '1'),
       };
     }
-
     return {
-      decision: SamplingDecision.RECORD_AND_SAMPLED,
+      ...wrapSamplingDecision({ decision: SamplingDecision.RECORD_AND_SAMPLED, context, spanAttributes }),
       attributes,
-      traceState,
     };
   }
 
@@ -152,3 +141,43 @@ function getParentSampled(parentSpan: Span, traceId: string, spanName: string):
 
   return undefined;
 }
+
+/**
+ * Wrap a sampling decision with data that Sentry needs to work properly with it.
+ * If you pass `decision: undefined`, it will be treated as `NOT_RECORDING`, but in contrast to passing `NOT_RECORDING`
+ * it will not propagate this decision to downstream Sentry SDKs.
+ */
+export function wrapSamplingDecision({
+  decision,
+  context,
+  spanAttributes,
+}: { decision: SamplingDecision | undefined; context: Context; spanAttributes: SpanAttributes }): SamplingResult {
+  const traceState = getBaseTraceState(context, spanAttributes);
+
+  // If the decision is undefined, we treat it as NOT_RECORDING, but we don't propagate this decision to downstream SDKs
+  // Which is done by not setting `SENTRY_TRACE_STATE_SAMPLED_NOT_RECORDING` traceState
+  if (decision == undefined) {
+    return { decision: SamplingDecision.NOT_RECORD, traceState };
+  }
+
+  if (decision === SamplingDecision.NOT_RECORD) {
+    return { decision, traceState: traceState.set(SENTRY_TRACE_STATE_SAMPLED_NOT_RECORDING, '1') };
+  }
+
+  return { decision, traceState };
+}
+
+function getBaseTraceState(context: Context, spanAttributes: SpanAttributes): TraceStateInterface {
+  const parentSpan = trace.getSpan(context);
+  const parentContext = parentSpan?.spanContext();
+
+  let traceState = parentContext?.traceState || new TraceState();
+
+  // We always keep the URL on the trace state, so we can access it in the propagator
+  const url = spanAttributes[SEMATTRS_HTTP_URL];
+  if (url && typeof url === 'string') {
+    traceState = traceState.set(SENTRY_TRACE_STATE_URL, url);
+  }
+
+  return traceState;
+}