feat(opentelemetry): Expose sampling helper

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;
+}