RUM-15138: Fix _dd.agent_psr to report the effective sample rate matching the actual sampling decision

applyPriority() and RequestTracingState.sampleRate now use a span-aware effectiveSampleRate() lookup instead of the argument-less getSampleRate(). For spans linked to an active RUM session the rebased rate is reported; for spans without a session ID the raw trace rate is reported, eliminating the mismatch where the backend received a rebased PSR for non-session spans sampled at the raw trace rate.

DeterministicTraceSampler gains an internal getSampleRate(DatadogSpan) method that resolves the effective rate per span, with a compatibility fallback to the session_sample_rate span tag for integrations that do not use ApmNetworkInstrumentationConfiguration. RumContextKeys centralises the session_sample_rate string constant across the trace module.

Author: hamorillo

features/dd-sdk-android-rum/src/test/kotlin/com/datadog/android/rum/internal/domain/RumContextTest.kt

@@ -7,6 +7,7 @@
 package com.datadog.android.rum.internal.domain
 
 import com.datadog.android.rum.utils.forge.Configurator
+import fr.xgouchet.elmyr.Forge
 import fr.xgouchet.elmyr.annotation.Forgery
 import fr.xgouchet.elmyr.junit5.ForgeConfiguration
 import fr.xgouchet.elmyr.junit5.ForgeExtension
@@ -31,16 +32,17 @@ internal class RumContextTest {
     }
 
     @Test
-    fun `M parse session sample rate W fromFeatureContext() {value is Double}`() {
-        // Given
+    fun `M parse session sample rate W fromFeatureContext() {value is Double}`(forge: Forge) {
+        // Given: use an explicit Double to verify that as? Number handles Double (as? Float would return null)
+        val fakeDouble: Double = forge.aDouble(min = 0.0, max = 100.0)
         val featureContext = mapOf<String, Any?>(
-            RumContext.SESSION_SAMPLE_RATE to 42.5
+            RumContext.SESSION_SAMPLE_RATE to fakeDouble
         )
 
         // When
         val rumContext = RumContext.fromFeatureContext(featureContext)
 
         // Then
-        assertThat(rumContext.sessionSampleRate).isEqualTo(42.5f)
+        assertThat(rumContext.sessionSampleRate).isEqualTo(fakeDouble.toFloat())
     }
 }

features/dd-sdk-android-rum/src/test/kotlin/com/datadog/android/rum/internal/domain/scope/RumSessionScopeTest.kt

@@ -493,19 +493,23 @@ internal class RumSessionScopeTest {
         forge: Forge
     ) {
         // Given
-        whenever(mockSessionSampler.getSampleRate()).thenReturn(100f, 20f)
+        // The scope field initializer consumes the first getSampleRate() call (irrelevant to the
+        // snapshot). The second call, inside renewSession(), is the one that gets snapshotted.
+        val fakeSessionRate = forge.aFloat(min = 0f, max = 100f)
+        val fakeLaterRate = forge.aFloat(min = 0f, max = 100f)
+        whenever(mockSessionSampler.getSampleRate()).thenReturn(RumContext.FULL_SESSION_SAMPLE_RATE, fakeSessionRate)
         whenever(mockSessionSampler.sample(any())).thenReturn(true)
         initializeTestedScope()
 
         // When
         testedScope.handleEvent(forge.startViewEvent(), fakeDatadogContext, mockEventWriteScope, mockWriter)
         val firstContext = testedScope.getRumContext()
-        whenever(mockSessionSampler.getSampleRate()).thenReturn(100f)
+        whenever(mockSessionSampler.getSampleRate()).thenReturn(fakeLaterRate)
         val secondContext = testedScope.getRumContext()
 
-        // Then
-        assertThat(firstContext.sessionSampleRate).isEqualTo(20f)
-        assertThat(secondContext.sessionSampleRate).isEqualTo(20f)
+        // Then — both contexts must report the rate snapshotted at session creation, not the changed rate
+        assertThat(firstContext.sessionSampleRate).isEqualTo(fakeSessionRate)
+        assertThat(secondContext.sessionSampleRate).isEqualTo(fakeSessionRate)
     }
 
     @Test

features/dd-sdk-android-trace/api/apiSurface

@@ -22,6 +22,7 @@ open class com.datadog.android.trace.DeterministicTraceSampler : com.datadog.and
   constructor(() -> Float)
   constructor(Float)
   constructor(Double)
+  override fun getSampleRate(): Float
   override fun sample(com.datadog.android.trace.api.span.DatadogSpan): Boolean
 annotation com.datadog.android.trace.ExperimentalTraceApi
 object com.datadog.android.trace.GlobalDatadogTracer

features/dd-sdk-android-trace/api/dd-sdk-android-trace.api

@@ -32,6 +32,7 @@ public class com/datadog/android/trace/DeterministicTraceSampler : com/datadog/a
 	public fun <init> (D)V
 	public fun <init> (F)V
 	public fun <init> (Lkotlin/jvm/functions/Function0;)V
+	public fun getSampleRate ()Ljava/lang/Float;
 	public fun sample (Lcom/datadog/android/trace/api/span/DatadogSpan;)Z
 	public synthetic fun sample (Ljava/lang/Object;)Z
 }

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/ApmNetworkInstrumentationConfiguration.kt

@@ -7,12 +7,16 @@ package com.datadog.android.trace
 
 import androidx.annotation.FloatRange
 import com.datadog.android.api.SdkCore
+import com.datadog.android.api.feature.Feature
+import com.datadog.android.core.InternalSdkCore
+import com.datadog.android.core.SdkReference
 import com.datadog.android.core.configuration.HostsSanitizer
 import com.datadog.android.core.internal.net.DefaultFirstPartyHostHeaderTypeResolver
 import com.datadog.android.core.sampling.Sampler
 import com.datadog.android.trace.api.span.DatadogSpan
 import com.datadog.android.trace.api.tracer.DatadogTracer
 import com.datadog.android.trace.internal.ApmNetworkInstrumentation
+import com.datadog.android.trace.internal.RumContextKeys
 import com.datadog.android.trace.internal.net.TracerProvider
 
 /**
@@ -46,7 +50,8 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
     internal var localTracerFactory: (SdkCore, Set<TracingHeaderType>) -> DatadogTracer = DEFAULT_LOCAL_TRACER_FACTORY,
     internal var traceContextInjection: TraceContextInjection = TraceContextInjection.SAMPLED,
     internal var tracedRequestListener: NetworkTracedRequestListener = NoOpNetworkTracedRequestListener(),
-    internal var traceSampler: Sampler<DatadogSpan> = DeterministicTraceSampler(DEFAULT_TRACE_SAMPLE_RATE),
+    internal var traceSampleRate: Float = DEFAULT_TRACE_SAMPLE_RATE,
+    internal var customTraceSampler: Sampler<DatadogSpan>? = null,
     internal var globalTracerProvider: () -> DatadogTracer? = { GlobalDatadogTracer.getOrNull() },
     internal var networkTracingScope: ApmNetworkTracingScope = ApmNetworkTracingScope.EXCLUDE_INTERNAL_REDIRECTS,
     internal var headerPropagationOnly: Boolean = false
@@ -116,7 +121,7 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
      * @param sampleRate the sample rate to use (percentage between 0f and 100f, default is 100f).
      */
     fun setTraceSampleRate(@FloatRange(from = 0.0, to = 100.0) sampleRate: Float) = apply {
-        this.traceSampler = DeterministicTraceSampler(sampleRate)
+        this.traceSampleRate = sampleRate
     }
 
     /**
@@ -127,7 +132,7 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
      * By default it is a sampler accepting 100% of the traces.
      */
     fun setTraceSampler(traceSampler: Sampler<DatadogSpan>) = apply {
-        this.traceSampler = traceSampler
+        this.customTraceSampler = traceSampler
     }
 
     /**
@@ -213,7 +218,8 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
         localTracerFactory = localTracerFactory,
         traceContextInjection = traceContextInjection,
         tracedRequestListener = tracedRequestListener,
-        traceSampler = traceSampler,
+        traceSampleRate = traceSampleRate,
+        customTraceSampler = customTraceSampler,
         globalTracerProvider = globalTracerProvider,
         networkTracingScope = networkTracingScope,
         headerPropagationOnly = headerPropagationOnly
@@ -241,10 +247,21 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
 
             val tracerProvider = TracerProvider(localTracerFactory, globalTracerProvider)
 
+            val sdkRef = SdkReference(sdkInstanceName)
+            val sessionSampleRateProvider: () -> Float = {
+                resolveSessionSampleRate(
+                    (sdkRef.get() as? InternalSdkCore)
+                        ?.getFeatureContext(Feature.RUM_FEATURE_NAME, useContextThread = false)
+                        ?.get(RumContextKeys.SESSION_SAMPLE_RATE)
+                )
+            }
+            val effectiveSampler = customTraceSampler
+                ?: DeterministicTraceSampler(traceSampleRate, sessionSampleRateProvider)
+
             return ApmNetworkInstrumentation(
                 canSendSpan = !headerPropagationOnly,
                 traceOrigin = traceOrigin,
-                traceSampler = traceSampler,
+                traceSampler = effectiveSampler,
                 tracerProvider = tracerProvider,
                 sdkInstanceName = sdkInstanceName,
                 injectionType = traceContextInjection,
@@ -270,6 +287,10 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
 
         private fun Map<String, Set<TracingHeaderType>>.deepCopy() = mapValues { (_, v) -> v.toSet() }
 
+        private fun resolveSessionSampleRate(rawSessionSampleRate: Any?): Float {
+            return (rawSessionSampleRate as? Number)?.toFloat() ?: DEFAULT_TRACE_SAMPLE_RATE
+        }
+
         private val DEFAULT_LOCAL_TRACER_FACTORY: (SdkCore, Set<TracingHeaderType>) -> DatadogTracer =
             { sdkCore, tracingHeaderTypes: Set<TracingHeaderType> ->
                 DatadogTracing.newTracerBuilder(sdkCore)

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/DeterministicTraceSampler.kt

@@ -8,35 +8,47 @@ package com.datadog.android.trace
 
 import androidx.annotation.FloatRange
 import com.datadog.android.core.sampling.DeterministicSampler
-import com.datadog.android.log.LogAttributes
 import com.datadog.android.trace.api.span.DatadogSpan
-import com.datadog.android.trace.internal.RumContextPropagator
+import com.datadog.android.trace.internal.RumContextKeys
 import com.datadog.android.trace.internal.net.SpanSamplingIdProvider
 
 /**
  * A [com.datadog.android.core.sampling.DeterministicSampler] using the TraceID of a Span to compute the sampling decision.
  *
- * When a span is linked to an active RUM session (i.e. the span carries a `rum.session.id` tag),
- * the sampling threshold is automatically rebased against the RUM session sample rate so that
- * the combined effective sampling probability reflects the configured trace sample rate applied
- * only to the sessions already tracked by RUM:
+ * The sampling threshold is automatically rebased against the RUM session sample rate when it is
+ * available so that the combined effective sampling probability reflects:
  *
  * ```
  * rebasedTraceSampleRate = sessionSampleRate * traceSampleRate / 100
  * ```
  *
- * @param sampleRateProvider Provider for the sample rate value which will be called each time
- * the sampling decision needs to be made. All the values should be in the range [0;100].
+ * Note: overriding [getSampleRate] is not a supported extension point for per-span sampling
+ * behavior.
+ *
+ * @param sampleRateProvider Provider for the trace sample rate. Called each time the sampling
+ * decision needs to be made. Values must be in the range [0;100].
+ * @param sessionSampleRateProvider Provider for the RUM session sample rate used to rebase the
+ * trace sample rate. Defaults to 100 (no rebasing).
  */
-open class DeterministicTraceSampler(
-    sampleRateProvider: () -> Float
+open class DeterministicTraceSampler private constructor(
+    private val sampleRateProvider: () -> Float,
+    private val sessionSampleRateProvider: () -> Float
 ) : DeterministicSampler<DatadogSpan>(
     SpanSamplingIdProvider::provideId,
     sampleRateProvider
 ) {
 
     /**
-     * Creates a new instance of [DeterministicSampler] with the given sample rate.
+     * Creates a new instance of [DeterministicTraceSampler] with the given sample rate provider.
+     *
+     * @param sampleRateProvider Provider for the sample rate value.
+     */
+    constructor(
+        sampleRateProvider: () -> Float
+    ) : this(sampleRateProvider, { DeterministicSampler.SAMPLE_ALL_RATE })
+
+    /**
+     * Creates a new instance of [DeterministicTraceSampler] with the given sample rate.
      *
      * @param sampleRate Sample rate to use.
      */
@@ -45,17 +57,48 @@ open class DeterministicTraceSampler(
     ) : this({ sampleRate })
 
     /**
-     * Creates a new instance of [DeterministicSampler] with the given sample rate.
+     * Creates a new instance of [DeterministicTraceSampler] with the given sample rate.
      *
      * @param sampleRate Sample rate to use.
      */
     constructor(
         @FloatRange(from = 0.0, to = 100.0) sampleRate: Double
     ) : this(sampleRate.toFloat())
 
+    /**
+     * Internal constructor used by [com.datadog.android.trace.ApmNetworkInstrumentationConfiguration]
+     * to wire the RUM session sample rate provider at instrumentation setup time.
+     */
+    internal constructor(
+        @FloatRange(from = 0.0, to = 100.0) sampleRate: Float,
+        sessionSampleRateProvider: () -> Float
+    ) : this({ sampleRate }, sessionSampleRateProvider)
+
+    /**
+     * Returns the effective sample rate, rebased against the RUM session sample rate when it is
+     * available:
+     *
+     * ```
+     * effectiveRate = traceSampleRate * sessionSampleRate / 100
+     * ```
+     *
+     * Falls back to the raw trace sample rate when no session sample rate is available
+     * (i.e. [sessionSampleRateProvider] returns 100).
+     */
+    override fun getSampleRate(): Float {
+        val traceSampleRate = super.getSampleRate()
+        val sessionSampleRate = sessionSampleRateProvider()
+        return if (sessionSampleRate >= 0f && sessionSampleRate < DeterministicSampler.SAMPLE_ALL_RATE) {
+            val rebased = traceSampleRate * sessionSampleRate / DeterministicSampler.SAMPLE_ALL_RATE
+            rebased.coerceAtMost(DeterministicSampler.SAMPLE_ALL_RATE)
+        } else {
+            traceSampleRate
+        }
+    }
+
     /** @inheritDoc */
     override fun sample(item: DatadogSpan): Boolean {
-        val sampleRate = rebasedSampleRate(item)
+        val sampleRate = getSampleRate(item)
         return when {
             sampleRate >= DeterministicSampler.SAMPLE_ALL_RATE -> true
             sampleRate <= 0f -> false
@@ -69,16 +112,25 @@ open class DeterministicTraceSampler(
         }
     }
 
-    private fun rebasedSampleRate(item: DatadogSpan): Float {
-        val traceSampleRate = getSampleRate()
-        val sessionId = item.context().tags[LogAttributes.RUM_SESSION_ID] as? String
-        val sessionSampleRate = item.context().tags[RumContextPropagator.SESSION_SAMPLE_RATE_KEY] as? Number
+    internal fun getSampleRate(item: DatadogSpan): Float {
+        // Compatibility path for integrations creating this sampler via public constructors
+        // (no sessionSampleRateProvider). In that case the session sample rate is propagated on
+        // the span context as `session_sample_rate` and must be used for rebasing.
+        val providerRate = sessionSampleRateProvider()
+        val providerRateIsValid = providerRate >= 0f && providerRate < DeterministicSampler.SAMPLE_ALL_RATE
+        val sessionSampleRate = when {
+            providerRateIsValid -> providerRate
+            else -> (item.context().tags[RumContextKeys.SESSION_SAMPLE_RATE] as? Number)?.toFloat()
+        }
 
-        return if (sessionId != null && sessionSampleRate != null) {
-            (traceSampleRate * sessionSampleRate.toFloat() / DeterministicSampler.SAMPLE_ALL_RATE)
-                .coerceAtMost(DeterministicSampler.SAMPLE_ALL_RATE)
+        return if (sessionSampleRate != null &&
+            sessionSampleRate >= 0f &&
+            sessionSampleRate < DeterministicSampler.SAMPLE_ALL_RATE
+        ) {
+            val rebased = super.getSampleRate() * sessionSampleRate / DeterministicSampler.SAMPLE_ALL_RATE
+            if (rebased > DeterministicSampler.SAMPLE_ALL_RATE) DeterministicSampler.SAMPLE_ALL_RATE else rebased
         } else {
-            traceSampleRate
+            super.getSampleRate()
         }
     }
 }

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/internal/ApmNetworkInstrumentation.kt

@@ -31,6 +31,7 @@ import com.datadog.android.trace.internal.net.RequestTracingState
 import com.datadog.android.trace.internal.net.TracerProvider
 import com.datadog.android.trace.internal.net.applyPriority
 import com.datadog.android.trace.internal.net.buildSpan
+import com.datadog.android.trace.internal.net.effectiveSampleRate
 import com.datadog.android.trace.internal.net.finishRumAware
 import com.datadog.android.trace.internal.net.sample
 import java.net.HttpURLConnection
@@ -143,7 +144,7 @@ class ApmNetworkInstrumentation internal constructor(
         return RequestTracingState(
             span = span,
             isSampled = isSampled,
-            sampleRate = traceSampler.getSampleRate(),
+            sampleRate = traceSampler.effectiveSampleRate(span),
             tracedRequestInfoBuilder = tracedRequestInfoBuilder
         )
     }

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/internal/RumContextKeys.kt

@@ -0,0 +1,12 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2016-Present Datadog, Inc.
+ */
+
+package com.datadog.android.trace.internal
+
+internal object RumContextKeys {
+    // Must match RumContext.SESSION_SAMPLE_RATE in dd-sdk-android-rum.
+    internal const val SESSION_SAMPLE_RATE = "session_sample_rate"
+}

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/internal/RumContextPropagator.kt

@@ -62,7 +62,10 @@ class RumContextPropagator(private val sdkCoreProvider: () -> FeatureSdkCore?) {
                 instance.setTag(LogAttributes.RUM_SESSION_ID, rumContext["session_id"])
                 instance.setTag(LogAttributes.RUM_VIEW_ID, rumContext["view_id"])
                 instance.setTag(LogAttributes.RUM_ACTION_ID, rumContext["action_id"])
-                instance.setTag(SESSION_SAMPLE_RATE_KEY, rumContext[SESSION_SAMPLE_RATE_KEY])
+                instance.setTag(
+                    RumContextKeys.SESSION_SAMPLE_RATE,
+                    rumContext[RumContextKeys.SESSION_SAMPLE_RATE]
+                )
                 instance.setTag(HttpCodec.RUM_KEY_USER_ID, datadogContext.userInfo.id)
                 instance.setTag(HttpCodec.RUM_KEY_ACCOUNT_ID, datadogContext.accountInfo?.id)
             }
@@ -110,7 +113,6 @@ class RumContextPropagator(private val sdkCoreProvider: () -> FeatureSdkCore?) {
 
     companion object {
         internal const val DATADOG_INITIAL_CONTEXT: String = "_dd.datadog_initial_context"
-        internal const val SESSION_SAMPLE_RATE_KEY: String = "session_sample_rate"
 
         internal const val INITIAL_DATADOG_CONTEXT_NOT_AVAILABLE_ERROR = "Initial span creation Datadog context" +
             " is not available at the write time."

features/dd-sdk-android-trace/src/main/kotlin/com/datadog/android/trace/internal/net/ApmNetworkInstrumentationExt.kt

@@ -9,6 +9,7 @@ import com.datadog.android.api.feature.Feature
 import com.datadog.android.api.feature.FeatureSdkCore
 import com.datadog.android.api.instrumentation.network.HttpRequestInfo
 import com.datadog.android.core.sampling.Sampler
+import com.datadog.android.trace.DeterministicTraceSampler
 import com.datadog.android.trace.api.DatadogTracingConstants.PrioritySampling
 import com.datadog.android.trace.api.DatadogTracingConstants.Tags
 import com.datadog.android.trace.api.span.DatadogSpan
@@ -35,11 +36,18 @@ internal fun DatadogSpan.applyPriority(isSampled: Boolean, traceSampler: Sampler
     if (spanContext.setSamplingPriority(samplingPriority)) {
         spanContext.setMetric(
             AGENT_PSR_ATTRIBUTE,
-            (traceSampler.getSampleRate() ?: ZERO_SAMPLE_RATE) / ALL_IN_SAMPLE_RATE
+            (traceSampler.effectiveSampleRate(this) ?: ZERO_SAMPLE_RATE) / ALL_IN_SAMPLE_RATE
         )
     }
 }
 
+internal fun Sampler<DatadogSpan>.effectiveSampleRate(span: DatadogSpan): Float? {
+    return when (this) {
+        is DeterministicTraceSampler -> getSampleRate(span)
+        else -> getSampleRate()
+    }
+}
+
 internal fun DatadogSpan.sample(request: HttpRequestInfo, traceSampler: Sampler<DatadogSpan>): Boolean {
     val samplingPriority = samplingPriority
     return if (samplingPriority != null) {