RUM-15138: Move session sample rate receiver ownership to TracingFeature

Instead of each ApmNetworkInstrumentation instance creating and managing its own FeatureContextUpdateReceiver (with manual close()/synchronized lifecycle), TracingFeature now owns a single receiver that listens for RUM context changes and broadcasts the session sample rate to all registered AtomicReference<Float> instances.

ApmNetworkInstrumentation sends a SessionSampleRateRegistrationEvent to TracingFeature on SDK resolution, which registers the ref and initializes it with the current rate. TracingFeature unregisters the receiver and clears all refs on onStop(), tying the lifecycle to the feature itself.

This eliminates the need for close() on ApmNetworkInstrumentation, fixing the leak where the OkHttp integration path never called close() while the Cronet path did.

# Conflicts:
#	features/dd-sdk-android-trace/src/test/kotlin/com/datadog/android/trace/internal/net/ApmNetworkInstrumentationTest.kt

Author: hamorillo

detekt_custom_safe_calls.yml

@@ -597,6 +597,7 @@ datadog:
       - "java.util.concurrent.atomic.AtomicReference.constructor(com.datadog.android.flags.model.FlagsClientState?)"
       - "java.util.concurrent.atomic.AtomicReference.constructor(com.datadog.android.flags.model.ProviderContext?)"
       - "java.util.concurrent.atomic.AtomicReference.constructor(com.datadog.android.rum.internal.domain.RumContext?)"
+      - "java.util.concurrent.atomic.AtomicReference.constructor(kotlin.Float?)"
       - "java.util.concurrent.atomic.AtomicReference.constructor(kotlin.collections.Map?)"
       - "java.util.concurrent.atomic.AtomicReference.constructor(kotlin.String?)"
       - "java.util.concurrent.atomic.AtomicReference.constructor(kotlin.collections.Set?)"
@@ -613,6 +614,7 @@ datadog:
       - "java.util.concurrent.atomic.AtomicReference.set(com.datadog.android.flags.model.ProviderContext?)"
       - "java.util.concurrent.atomic.AtomicReference.set(com.datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI?)"
       - "java.util.concurrent.atomic.AtomicReference.set(com.datadog.trace.core.CoreTracer?)"
+      - "java.util.concurrent.atomic.AtomicReference.set(kotlin.Float?)"
       - "java.util.concurrent.atomic.AtomicReference.set(kotlin.collections.Map?)"
       - "java.util.concurrent.atomic.AtomicReference.set(kotlin.Nothing?)"
       - "java.util.concurrent.atomic.AtomicReference.set(kotlin.String?)"
@@ -997,6 +999,7 @@ datadog:
       - "kotlin.collections.MutableSet.add(com.datadog.android.sessionreplay.ExtensionSupport)"
       - "kotlin.collections.MutableSet.add(com.datadog.android.telemetry.internal.TelemetryEventId)"
       - "kotlin.collections.MutableSet.add(java.io.File)"
+      - "kotlin.collections.MutableSet.add(java.util.concurrent.atomic.AtomicReference)"
       - "kotlin.collections.MutableSet.add(kotlin.String)"
       - "kotlin.collections.MutableSet.add(kotlin.String?)"
       - "kotlin.collections.MutableSet.addAll(kotlin.collections.Collection)"

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

@@ -59,7 +59,6 @@ class com.datadog.android.trace.internal.ApmNetworkInstrumentation
   fun onResponseFailed(com.datadog.android.trace.internal.net.RequestTracingState, Throwable)
   fun removeTracingHeaders(com.datadog.android.api.instrumentation.network.HttpRequestInfoBuilder)
   fun reportInstrumentationError(() -> String)
-  fun close()
 class com.datadog.android.trace.internal.DatadogPropagationHelper
   fun isExtractedContext(com.datadog.android.trace.api.span.DatadogSpanContext): Boolean
   fun setTraceContext(com.datadog.android.api.instrumentation.network.HttpRequestInfoBuilder, String, String, Int)

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

@@ -99,7 +99,6 @@ public abstract interface class com/datadog/android/trace/event/SpanEventMapper
 }
 
 public final class com/datadog/android/trace/internal/ApmNetworkInstrumentation {
-	public final fun close ()V
 	public final fun getLocalHeaderTypes ()Ljava/util/Set;
 	public final fun getNetworkTracingScope ()Lcom/datadog/android/trace/ApmNetworkTracingScope;
 	public final fun getSampleRate ()Ljava/lang/Float;

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

@@ -7,15 +7,13 @@ 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.api.feature.FeatureContextUpdateReceiver
 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.log.LogAttributes
 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.TracingFeature
 import com.datadog.android.trace.internal.net.TracerProvider
 import java.util.concurrent.atomic.AtomicReference
 
@@ -242,11 +240,6 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
         internal const val ALL_IN_SAMPLE_RATE: Double = 100.0
         private const val DEFAULT_TRACE_SAMPLE_RATE: Float = 100f
 
-        // Sentinel value meaning "assume all RUM sessions are sampled" so that no session-based
-        // rebasing is applied until the RUM feature publishes its actual session sample rate.
-        // With this value the effective rate is: traceSampleRate * 100 / 100 = traceSampleRate.
-        private const val NO_SESSION_REBASING_RATE: Float = 100f
-
         internal const val NETWORK_REQUESTS_TRACKING_FEATURE_NAME = "Network Requests"
 
         internal fun ApmNetworkInstrumentationConfiguration.createInstrumentation(
@@ -259,22 +252,9 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
             val tracerProvider = TracerProvider(localTracerFactory, globalTracerProvider)
 
             @Suppress("UnsafeThirdPartyFunctionCall") // AtomicReference initialized with non-null Float constant
-            val cachedSessionSampleRate = AtomicReference(NO_SESSION_REBASING_RATE)
-            val sessionSampleRateReceiver: FeatureContextUpdateReceiver? =
-                if (customTraceSampler == null) {
-                    FeatureContextUpdateReceiver { featureName, context ->
-                        if (featureName == Feature.RUM_FEATURE_NAME) {
-                            @Suppress(
-                                "UnsafeThirdPartyFunctionCall"
-                            ) // resolveSessionSampleRate always returns non-null Float
-                            cachedSessionSampleRate.set(
-                                resolveSessionSampleRate(context[LogAttributes.RUM_SESSION_SAMPLE_RATE])
-                            )
-                        }
-                    }
-                } else {
-                    null
-                }
+            val cachedSessionSampleRate = AtomicReference(TracingFeature.NO_SESSION_REBASING_RATE)
+            val sessionSampleRateRef: AtomicReference<Float>? =
+                if (customTraceSampler == null) cachedSessionSampleRate else null
             val effectiveSampler = customTraceSampler
                 ?: DeterministicTraceSampler(traceSampleRate) { cachedSessionSampleRate.get() }
 
@@ -290,7 +270,7 @@ class ApmNetworkInstrumentationConfiguration internal constructor(
                 tracedRequestListener = tracedRequestListener,
                 redacted404ResourceName = redacted404ResourceName,
                 localFirstPartyHostHeaderTypeResolver = localFirstPartyHostHeaderTypeResolver,
-                sessionSampleRateReceiver = sessionSampleRateReceiver
+                sessionSampleRateRef = sessionSampleRateRef
             )
         }
 
@@ -308,10 +288,6 @@ 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() ?: NO_SESSION_REBASING_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/internal/ApmNetworkInstrumentation.kt

@@ -8,7 +8,7 @@ package com.datadog.android.trace.internal
 
 import com.datadog.android.api.InternalLogger
 import com.datadog.android.api.SdkCore
-import com.datadog.android.api.feature.FeatureContextUpdateReceiver
+import com.datadog.android.api.feature.Feature
 import com.datadog.android.api.instrumentation.network.HttpRequestInfo
 import com.datadog.android.api.instrumentation.network.HttpRequestInfoBuilder
 import com.datadog.android.api.instrumentation.network.HttpResponseInfo
@@ -38,6 +38,7 @@ import com.datadog.android.trace.internal.net.finishRumAware
 import com.datadog.android.trace.internal.net.sample
 import java.net.HttpURLConnection
 import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
 
 /**
  * For internal usage only.
@@ -58,7 +59,11 @@ import java.util.Locale
  * @param localFirstPartyHostHeaderTypeResolver resolver for determining header types for first-party hosts.
  * @param networkingLibraryName the name identifying the network instrumentation (e.g., "OkHttp", "Cronet").
  * @param networkTracingScope Tracing scope for the instrumentation. See [ApmNetworkTracingScope] enum for more details.
- * @param sessionSampleRateReceiver optional receiver that caches the RUM session sample rate for sampling decisions. When provided, it is registered with the SDK on first resolution and updated whenever the RUM context changes.
+ * @param sessionSampleRateRef optional [AtomicReference] that [TracingFeature] will keep
+ *        up-to-date with the latest RUM session sample rate. When the SDK core resolves,
+ *        this ref is sent to [TracingFeature] via `sendEvent` so that it is updated whenever
+ *        the RUM context changes. The corresponding [DeterministicTraceSampler] reads from
+ *        this ref on every request (volatile read, no locking).
  */
 @Suppress("LongParameterList")
 @InternalApi
@@ -74,12 +79,9 @@ class ApmNetworkInstrumentation internal constructor(
     internal val localFirstPartyHostHeaderTypeResolver: DefaultFirstPartyHostHeaderTypeResolver,
     private val networkingLibraryName: String,
     val networkTracingScope: ApmNetworkTracingScope = ApmNetworkTracingScope.ALL,
-    internal val sessionSampleRateReceiver: FeatureContextUpdateReceiver? = null
+    internal val sessionSampleRateRef: AtomicReference<Float>? = null
 ) {
     private val rumContextPropagator = RumContextPropagator { internalSdkCore }
-    private val receiverLifecycleLock = Any()
-    private var isClosed = false
-    private var isRegistered = false
 
     private val internalSdkCore: InternalSdkCore?
         get() = sdkCoreReference.get() as? InternalSdkCore
@@ -96,13 +98,9 @@ class ApmNetworkInstrumentation internal constructor(
                 WARNING_TRACING_NO_HOSTS.format(Locale.US, networkingLibraryName)
             }
         }
-        sessionSampleRateReceiver?.let { receiver ->
-            synchronized(receiverLifecycleLock) {
-                if (!isClosed && !isRegistered) {
-                    sdkCore.setContextUpdateReceiver(receiver)
-                    isRegistered = true
-                }
-            }
+        sessionSampleRateRef?.let { ref ->
+            sdkCore.getFeature(Feature.TRACING_FEATURE_NAME)
+                ?.sendEvent(SessionSampleRateRegistrationEvent(ref))
         }
     }
 
@@ -234,20 +232,6 @@ class ApmNetworkInstrumentation internal constructor(
         )
     }
 
-    /**
-     * Releases this instrumentation and unregisters context receivers from SDK core.
-     */
-    fun close() {
-        val receiver = sessionSampleRateReceiver ?: return
-        synchronized(receiverLifecycleLock) {
-            isClosed = true
-            if (isRegistered) {
-                isRegistered = false
-                internalSdkCore?.removeContextUpdateReceiver(receiver)
-            }
-        }
-    }
-
     private fun RequestTracingState.onRequestIntercepted(
         response: HttpResponseInfo?,
         throwable: Throwable?

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

@@ -0,0 +1,23 @@
+/*
+ * 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
+
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * Event sent by [ApmNetworkInstrumentation] to [TracingFeature] via `sendEvent`
+ * to register an [AtomicReference] that should be kept up-to-date with the current
+ * RUM session sample rate.
+ *
+ * TracingFeature updates all registered references whenever the RUM feature context
+ * changes, ensuring that [com.datadog.android.trace.DeterministicTraceSampler] can
+ * read the latest session sample rate with a simple volatile read (no locking).
+ *
+ * @param ref the [AtomicReference] that TracingFeature will update with the
+ * latest RUM session sample rate.
+ */
+internal data class SessionSampleRateRegistrationEvent(val ref: AtomicReference<Float>)

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

@@ -8,10 +8,13 @@ package com.datadog.android.trace.internal
 
 import android.content.Context
 import com.datadog.android.api.feature.Feature
+import com.datadog.android.api.feature.FeatureContextUpdateReceiver
+import com.datadog.android.api.feature.FeatureEventReceiver
 import com.datadog.android.api.feature.FeatureSdkCore
 import com.datadog.android.api.feature.StorageBackedFeature
 import com.datadog.android.api.net.RequestFactory
 import com.datadog.android.api.storage.FeatureStorageConfiguration
+import com.datadog.android.log.LogAttributes
 import com.datadog.android.trace.InternalCoreWriterProvider
 import com.datadog.android.trace.event.SpanEventMapper
 import com.datadog.android.trace.internal.data.CoreTraceWriter
@@ -22,6 +25,7 @@ import com.datadog.android.trace.internal.net.TracesRequestFactory
 import com.datadog.trace.common.writer.NoOpWriter
 import com.datadog.trace.common.writer.Writer
 import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicReference
 
 /**
  * Tracing feature class, which needs to be registered with Datadog SDK instance.
@@ -31,17 +35,23 @@ internal class TracingFeature(
     customEndpointUrl: String?,
     internal val spanEventMapper: SpanEventMapper,
     internal val networkInfoEnabled: Boolean
-) : InternalCoreWriterProvider, StorageBackedFeature {
+) : InternalCoreWriterProvider, StorageBackedFeature, FeatureContextUpdateReceiver, FeatureEventReceiver {
 
     internal var coreTracerDataWriter: Writer = NoOpWriter()
     internal val initialized = AtomicBoolean(false)
 
+    internal val internalSessionSampleRate = AtomicReference(NO_SESSION_REBASING_RATE)
+    private val sessionSampleRateRefs = mutableSetOf<AtomicReference<Float>>()
+    internal val registrationLock = Any()
+
     // region Feature
 
     override val name: String = Feature.TRACING_FEATURE_NAME
 
     override fun onInitialize(appContext: Context) {
         coreTracerDataWriter = createDataWriter(sdkCore)
+        sdkCore.setContextUpdateReceiver(this)
+        sdkCore.setEventReceiver(name, this)
         initialized.set(true)
     }
 
@@ -56,17 +66,54 @@ internal class TracingFeature(
         FeatureStorageConfiguration.DEFAULT
 
     override fun onStop() {
+        sdkCore.removeContextUpdateReceiver(this)
+        sdkCore.removeEventReceiver(name)
+        resetSessionSampleRateTracking()
         initialized.set(false)
     }
 
     // endregion
 
+    // region FeatureContextUpdateReceiver
+
+    override fun onContextUpdate(featureName: String, context: Map<String, Any?>) {
+        if (featureName == Feature.RUM_FEATURE_NAME) {
+            val rate = resolveSessionSampleRate(context[LogAttributes.RUM_SESSION_SAMPLE_RATE])
+            synchronized(registrationLock) {
+                internalSessionSampleRate.set(rate)
+                sessionSampleRateRefs.forEach { it.set(rate) }
+            }
+        }
+    }
+
+    // endregion
+
+    // region FeatureEventReceiver
+
+    override fun onReceive(event: Any) {
+        if (event is SessionSampleRateRegistrationEvent) {
+            synchronized(registrationLock) {
+                sessionSampleRateRefs.add(event.ref)
+                event.ref.set(internalSessionSampleRate.get())
+            }
+        }
+    }
+
+    // endregion
+
     // region InternalCoreWriterProvider
 
     override fun getCoreTracerWriter() = DatadogSpanWriterWrapper(coreTracerDataWriter)
 
     // endregion
 
+    private fun resetSessionSampleRateTracking() {
+        synchronized(registrationLock) {
+            sessionSampleRateRefs.clear()
+            internalSessionSampleRate.set(NO_SESSION_REBASING_RATE)
+        }
+    }
+
     private fun createDataWriter(sdkCore: FeatureSdkCore): Writer {
         val internalLogger = sdkCore.internalLogger
         return CoreTraceWriter(
@@ -77,4 +124,15 @@ internal class TracingFeature(
             internalLogger = internalLogger
         )
     }
+
+    internal companion object {
+        // Sentinel value meaning "assume all RUM sessions are sampled" so that no session-based
+        // rebasing is applied until the RUM feature publishes its actual session sample rate.
+        // With this value the effective rate is: traceSampleRate * 100 / 100 = traceSampleRate.
+        internal const val NO_SESSION_REBASING_RATE: Float = 100f
+
+        private fun resolveSessionSampleRate(rawSessionSampleRate: Any?): Float {
+            return (rawSessionSampleRate as? Number)?.toFloat() ?: NO_SESSION_REBASING_RATE
+        }
+    }
 }