RUM-6866 Allow typed Sampler

DataDog · Nov 13, 2024 · 808f3d7 · 808f3d7
1 parent a7c4f31
commit 808f3d7
Show file tree

Hide file tree

Showing 26 changed files with 256 additions and 201 deletions.
diff --git a/dd-sdk-android-core/api/apiSurface b/dd-sdk-android-core/api/apiSurface
@@ -323,14 +323,14 @@ interface com.datadog.android.core.persistence.Serializer<T: Any>
 fun <T: Any> Serializer<T>.serializeToByteArray(T, com.datadog.android.api.InternalLogger): ByteArray?
 data class com.datadog.android.core.persistence.datastore.DataStoreContent<T: Any>
   constructor(Int, T?)
-class com.datadog.android.core.sampling.RateBasedSampler : Sampler
+open class com.datadog.android.core.sampling.RateBasedSampler<T: Any> : Sampler<T>
   constructor(() -> Float)
   constructor(Float)
   constructor(Double)
-  override fun sample(): Boolean
+  override fun sample(T): Boolean
   override fun getSampleRate(): Float
-interface com.datadog.android.core.sampling.Sampler
-  fun sample(): Boolean
+interface com.datadog.android.core.sampling.Sampler<T: Any>
+  fun sample(T): Boolean
   fun getSampleRate(): Float?
 interface com.datadog.android.core.thread.FlushableExecutorService : java.util.concurrent.ExecutorService
   fun drainTo(MutableCollection<Runnable>)

diff --git a/dd-sdk-android-core/api/dd-sdk-android-core.api b/dd-sdk-android-core/api/dd-sdk-android-core.api
@@ -867,18 +867,18 @@ public final class com/datadog/android/core/persistence/datastore/DataStoreConte
 	public fun toString ()Ljava/lang/String;
 }
 
-public final class com/datadog/android/core/sampling/RateBasedSampler : com/datadog/android/core/sampling/Sampler {
+public class com/datadog/android/core/sampling/RateBasedSampler : com/datadog/android/core/sampling/Sampler {
 	public static final field SAMPLE_ALL_RATE F
 	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 ()Z
+	public fun sample (Ljava/lang/Object;)Z
 }
 
 public abstract interface class com/datadog/android/core/sampling/Sampler {
 	public abstract fun getSampleRate ()Ljava/lang/Float;
-	public abstract fun sample ()Z
+	public abstract fun sample (Ljava/lang/Object;)Z
 }
 
 public abstract interface class com/datadog/android/core/thread/FlushableExecutorService : java/util/concurrent/ExecutorService {

diff --git a/...ndroid-core/src/main/kotlin/com/datadog/android/core/internal/logger/SdkInternalLogger.kt b/...ndroid-core/src/main/kotlin/com/datadog/android/core/internal/logger/SdkInternalLogger.kt
@@ -98,7 +98,7 @@ internal class SdkInternalLogger(
         additionalProperties: Map<String, Any?>,
         samplingRate: Float
     ) {
-        if (!RateBasedSampler(samplingRate).sample()) return
+        if (!sample(samplingRate)) return
         val rumFeature = sdkCore?.getFeature(Feature.RUM_FEATURE_NAME) ?: return
         val metricEvent = InternalTelemetryEvent.Metric(
             message = messageBuilder(),
@@ -113,7 +113,7 @@ internal class SdkInternalLogger(
         samplingRate: Float,
         operationName: String
     ): PerformanceMetric? {
-        if (!RateBasedSampler(samplingRate).sample()) return null
+        if (!sample(samplingRate)) return null
 
         return when (metric) {
             TelemetryMetricType.MethodCalled -> {
@@ -130,7 +130,7 @@ internal class SdkInternalLogger(
         apiUsageEvent: InternalTelemetryEvent.ApiUsage,
         samplingRate: Float
     ) {
-        if (!RateBasedSampler(samplingRate).sample()) return
+        if (!sample(samplingRate)) return
         val rumFeature = sdkCore?.getFeature(Feature.RUM_FEATURE_NAME) ?: return
         rumFeature.sendEvent(apiUsageEvent)
     }
@@ -139,6 +139,10 @@ internal class SdkInternalLogger(
 
     // region Internal
 
+    fun sample(samplingRate: Float): Boolean {
+        return RateBasedSampler<Unit>(samplingRate).sample(Unit)
+    }
+
     private fun logToUser(
         level: InternalLogger.Level,
         messageBuilder: () -> String,

diff --git a/dd-sdk-android-core/src/main/kotlin/com/datadog/android/core/sampling/RateBasedSampler.kt b/dd-sdk-android-core/src/main/kotlin/com/datadog/android/core/sampling/RateBasedSampler.kt
@@ -13,10 +13,11 @@ import java.security.SecureRandom
 /**
  * [Sampler] with the given sample rate which can be fixed or dynamic.
  *
+ * @param T the type of items to sample.
  * @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 on the scale [0;100].
  */
-class RateBasedSampler(private val sampleRateProvider: () -> Float) : Sampler {
+open class RateBasedSampler<T : Any>(private val sampleRateProvider: () -> Float) : Sampler<T> {
 
     /**
      * Creates a new instance of [RateBasedSampler] with the given sample rate.
@@ -36,7 +37,7 @@ class RateBasedSampler(private val sampleRateProvider: () -> Float) : Sampler {
 
     /** @inheritDoc */
     @Suppress("MagicNumber")
-    override fun sample(): Boolean {
+    override fun sample(item: T): Boolean {
         return when (val sampleRate = getSampleRate()) {
             0f -> false
             SAMPLE_ALL_RATE -> true

diff --git a/dd-sdk-android-core/src/main/kotlin/com/datadog/android/core/sampling/Sampler.kt b/dd-sdk-android-core/src/main/kotlin/com/datadog/android/core/sampling/Sampler.kt
@@ -10,14 +10,15 @@ import androidx.annotation.FloatRange
 
 /**
  * Interface representing the sampling.
+ * @param T the type of items to sample.
  */
-interface Sampler {
+interface Sampler<T : Any> {
 
     /**
-     * Sampling method.
-     * @return true if you want to keep the value, false otherwise.
+     * @param item the item to sample
+     * @return true to keep the item, false to discard it
      */
-    fun sample(): Boolean
+    fun sample(item: T): Boolean
 
     /**
      * @return the sample rate if applicable, as a float between 0 and 100,

diff --git a/...dk-android-core/src/test/kotlin/com/datadog/android/core/sampling/RateBasedSamplerTest.kt b/...dk-android-core/src/test/kotlin/com/datadog/android/core/sampling/RateBasedSamplerTest.kt
@@ -30,7 +30,7 @@ import kotlin.math.sqrt
 @ForgeConfiguration(Configurator::class)
 internal class RateBasedSamplerTest {
 
-    private lateinit var testedSampler: RateBasedSampler
+    private lateinit var testedSampler: RateBasedSampler<Unit>
 
     @FloatForgery(min = 0f, max = 100f)
     var randomSampleRate: Float = 0.0f
@@ -51,7 +51,7 @@ internal class RateBasedSamplerTest {
         repeat(testRepeats) {
             var validated = 0
             repeat(dataSize) {
-                val isValid = if (testedSampler.sample()) 1 else 0
+                val isValid = if (testedSampler.sample(Unit)) 1 else 0
                 validated += isValid
             }
             val computedSampleRate = (validated.toDouble() / dataSize.toDouble()) * 100
@@ -93,7 +93,7 @@ internal class RateBasedSamplerTest {
         repeat(testRepeats) {
             var validated = 0
             repeat(dataSize) {
-                val isValid = if (testedSampler.sample()) 1 else 0
+                val isValid = if (testedSampler.sample(Unit)) 1 else 0
                 validated += isValid
             }
             val computedSampleRate = (validated.toDouble() / dataSize.toDouble()) * 100
@@ -120,7 +120,7 @@ internal class RateBasedSamplerTest {
         val dataSize = 10
 
         repeat(dataSize) {
-            val isValid = if (testedSampler.sample()) 1 else 0
+            val isValid = if (testedSampler.sample(Unit)) 1 else 0
             validated += isValid
         }
 
@@ -135,7 +135,7 @@ internal class RateBasedSamplerTest {
         val dataSize = 10
 
         repeat(dataSize) {
-            val isValid = if (testedSampler.sample()) 1 else 0
+            val isValid = if (testedSampler.sample(Unit)) 1 else 0
             validated += isValid
         }
 

diff --git a/...android-logs/src/main/kotlin/com/datadog/android/log/internal/logger/DatadogLogHandler.kt b/...android-logs/src/main/kotlin/com/datadog/android/log/internal/logger/DatadogLogHandler.kt
@@ -27,7 +27,7 @@ internal class DatadogLogHandler(
     internal val attachNetworkInfo: Boolean,
     internal val bundleWithTraces: Boolean = true,
     internal val bundleWithRum: Boolean = true,
-    internal val sampler: Sampler = RateBasedSampler(DEFAULT_SAMPLE_RATE),
+    internal val sampler: Sampler<Unit> = RateBasedSampler(DEFAULT_SAMPLE_RATE),
     internal val minLogPriority: Int = -1
 ) : LogHandler {
 
@@ -52,7 +52,7 @@ internal class DatadogLogHandler(
             combinedAttributes.putAll(logsFeature.unwrap<LogsFeature>().getAttributes().toMutableMap())
         }
         combinedAttributes.putAll(attributes)
-        if (sampler.sample()) {
+        if (sampler.sample(Unit)) {
             if (logsFeature != null) {
                 val threadName = Thread.currentThread().name
                 logsFeature.withWriteContext { datadogContext, eventBatchWriter ->
@@ -123,7 +123,7 @@ internal class DatadogLogHandler(
         }
         combinedAttributes.putAll(attributes)
 
-        if (sampler.sample()) {
+        if (sampler.sample(Unit)) {
             if (logsFeature != null) {
                 val threadName = Thread.currentThread().name
                 logsFeature.withWriteContext { datadogContext, eventBatchWriter ->

diff --git a/...oid-logs/src/test/kotlin/com/datadog/android/log/internal/logger/DatadogLogHandlerTest.kt b/...oid-logs/src/test/kotlin/com/datadog/android/log/internal/logger/DatadogLogHandlerTest.kt
@@ -106,7 +106,7 @@ internal class DatadogLogHandlerTest {
     lateinit var mockWriter: DataWriter<LogEvent>
 
     @Mock
-    lateinit var mockSampler: Sampler
+    lateinit var mockSampler: Sampler<Unit>
 
     @BeforeEach
     fun `set up`(forge: Forge) {
@@ -937,7 +937,7 @@ internal class DatadogLogHandlerTest {
     @Test
     fun `it will sample out the logs when required`() {
         // Given
-        whenever(mockSampler.sample()).thenReturn(false)
+        whenever(mockSampler.sample(Unit)).thenReturn(false)
         testedHandler = DatadogLogHandler(
             loggerName = fakeLoggerName,
             logGenerator = DatadogLogGenerator(
@@ -967,7 +967,7 @@ internal class DatadogLogHandlerTest {
     fun `it will sample in the logs when required`() {
         // Given
         val now = System.currentTimeMillis()
-        whenever(mockSampler.sample()).thenReturn(true)
+        whenever(mockSampler.sample(Unit)).thenReturn(true)
         testedHandler = DatadogLogHandler(
             loggerName = fakeLoggerName,
             logGenerator = DatadogLogGenerator(

diff --git a/...droid-rum/src/main/kotlin/com/datadog/android/telemetry/internal/TelemetryEventHandler.kt b/...droid-rum/src/main/kotlin/com/datadog/android/telemetry/internal/TelemetryEventHandler.kt
@@ -35,8 +35,9 @@ import com.datadog.android.telemetry.model.TelemetryConfigurationEvent.ViewTrack
 @Suppress("TooManyFunctions")
 internal class TelemetryEventHandler(
     internal val sdkCore: InternalSdkCore,
-    internal val eventSampler: Sampler,
-    internal val configurationExtraSampler: Sampler = RateBasedSampler(DEFAULT_CONFIGURATION_SAMPLE_RATE),
+    internal val eventSampler: Sampler<InternalTelemetryEvent>,
+    internal val configurationExtraSampler: Sampler<InternalTelemetryEvent> =
+        RateBasedSampler(DEFAULT_CONFIGURATION_SAMPLE_RATE),
     private val sessionEndedMetricDispatcher: SessionMetricDispatcher,
     private val maxEventCountPerSession: Int = MAX_EVENTS_PER_SESSION
 ) : RumSessionListener {
@@ -125,13 +126,13 @@ internal class TelemetryEventHandler(
         totalEventsSeenInCurrentSession = 0
     }
 
-    // region private
+    // region Internal
 
     @Suppress("ReturnCount")
     private fun canWrite(event: InternalTelemetryEvent): Boolean {
-        if (!eventSampler.sample()) return false
+        if (!eventSampler.sample(event)) return false
 
-        if (event is InternalTelemetryEvent.Configuration && !configurationExtraSampler.sample()) {
+        if (event is InternalTelemetryEvent.Configuration && !configurationExtraSampler.sample(event)) {
             return false
         }
 

diff --git a/...d-rum/src/test/kotlin/com/datadog/android/telemetry/internal/TelemetryEventHandlerTest.kt b/...d-rum/src/test/kotlin/com/datadog/android/telemetry/internal/TelemetryEventHandlerTest.kt
@@ -90,10 +90,10 @@ internal class TelemetryEventHandlerTest {
     lateinit var mockWriter: DataWriter<Any>
 
     @Mock
-    lateinit var mockSampler: Sampler
+    lateinit var mockSampler: Sampler<InternalTelemetryEvent>
 
     @Mock
-    lateinit var mockConfigurationSampler: Sampler
+    lateinit var mockConfigurationSampler: Sampler<InternalTelemetryEvent>
 
     @Mock
     lateinit var mockSdkCore: InternalSdkCore
@@ -170,8 +170,8 @@ internal class TelemetryEventHandlerTest {
             deviceInfo = mockDeviceInfo
         )
 
-        whenever(mockSampler.sample()) doReturn true
-        whenever(mockConfigurationSampler.sample()) doReturn true
+        whenever(mockSampler.sample(any())) doReturn true
+        whenever(mockConfigurationSampler.sample(any())) doReturn true
 
         whenever(
             mockSdkCore.getFeature(Feature.RUM_FEATURE_NAME)
@@ -744,7 +744,7 @@ internal class TelemetryEventHandlerTest {
         // Given
         val fakeInternalTelemetryEvent = forge.forgeWritableInternalTelemetryEvent()
         val rawEvent = RumRawEvent.TelemetryEventWrapper(fakeInternalTelemetryEvent)
-        whenever(mockSampler.sample()) doReturn false
+        whenever(mockSampler.sample(any())) doReturn false
 
         // When
         testedTelemetryHandler.handleEvent(rawEvent, mockWriter)
@@ -763,8 +763,8 @@ internal class TelemetryEventHandlerTest {
             forge.getForgery<InternalTelemetryEvent.Log.Debug>()
         )
         val rawEvent = RumRawEvent.TelemetryEventWrapper(logeEvent)
-        whenever(mockSampler.sample()) doReturn true
-        whenever(mockConfigurationSampler.sample()) doReturn false
+        whenever(mockSampler.sample(any())) doReturn true
+        whenever(mockConfigurationSampler.sample(any())) doReturn false
 
         // When
         testedTelemetryHandler.handleEvent(rawEvent, mockWriter)
@@ -786,8 +786,8 @@ internal class TelemetryEventHandlerTest {
     ) {
         // Given
         val rawEvent = RumRawEvent.TelemetryEventWrapper(fakeConfiguration)
-        whenever(mockSampler.sample()) doReturn true
-        whenever(mockConfigurationSampler.sample()) doReturn false
+        whenever(mockSampler.sample(any())) doReturn true
+        whenever(mockConfigurationSampler.sample(any())) doReturn false
 
         // When
         testedTelemetryHandler.handleEvent(rawEvent, mockWriter)
@@ -930,7 +930,7 @@ internal class TelemetryEventHandlerTest {
     fun `M count the limit only after the sampling W handleEvent()`(forge: Forge) {
         // Given
         // sample out 50%
-        whenever(mockSampler.sample()) doAnswer object : Answer<Boolean> {
+        whenever(mockSampler.sample(any())) doAnswer object : Answer<Boolean> {
             var invocationCount = 0
             override fun answer(invocation: InvocationOnMock): Boolean {
                 invocationCount++
@@ -1008,8 +1008,8 @@ internal class TelemetryEventHandlerTest {
         }
     }
 
-// endregion
-// region Api Usage
+    // endregion
+    // region Api Usage
 
     @Test
     fun `M create api usage event W handleEvent(api usage event)`(
@@ -1081,9 +1081,10 @@ internal class TelemetryEventHandlerTest {
             message = TelemetryEventHandler.MAX_EVENT_NUMBER_REACHED_MESSAGE
         )
     }
-// endregion
 
-// region Assertions
+    // endregion
+
+    // region Assertions
 
     private fun assertApiUsageMatchesInternalEvent(
         actual: TelemetryUsageEvent,
@@ -1257,7 +1258,7 @@ internal class TelemetryEventHandlerTest {
         )
     }
 
-// endregion
+    // endregion
 
     companion object {
 

diff --git a/...replay/src/main/kotlin/com/datadog/android/sessionreplay/internal/SessionReplayFeature.kt b/...replay/src/main/kotlin/com/datadog/android/sessionreplay/internal/SessionReplayFeature.kt
@@ -50,7 +50,7 @@ internal class SessionReplayFeature(
     internal val textAndInputPrivacy: TextAndInputPrivacy,
     internal val touchPrivacy: TouchPrivacy,
     internal val imagePrivacy: ImagePrivacy,
-    private val rateBasedSampler: Sampler,
+    private val rateBasedSampler: Sampler<Unit>,
     private val startRecordingImmediately: Boolean,
     private val recorderProvider: RecorderProvider
 ) : StorageBackedFeature, FeatureEventReceiver {
@@ -261,7 +261,7 @@ internal class SessionReplayFeature(
 
     private fun applySampling(alreadySeenSession: Boolean) {
         if (!alreadySeenSession) {
-            isSessionSampledIn.set(rateBasedSampler.sample())
+            isSessionSampledIn.set(rateBasedSampler.sample(Unit))
         }
     }