Add session ID ratio based sampler (#262)

* Add  session ID ratio based sampler

* Remove redundant configuration option.

* Rename sampling config option and keep boolean field internally.

* Update session sampling option readme

* Remove parentBased around sampler

* profiler: Update documentation

Author: agoallikmaa

Diff for: README.md

@@ -256,7 +256,9 @@ when initializing your instance of the SplunkRum API:
 - `limitDiskUsageMegabytes(int)` :
   When disk buffering is enabled, this can be used to adjust the maximum amount of storage
   that will be used. Default = 25MB.
-
+- `enableSessionBasedSampling(double)` :
+  Enable session ID based sampling and set its sampling ratio. The ratio is a probability of a
+  session being included between between 0.0 (all dropped) and 1.0 (all included).
 
 #### APIs provided by the `SplunkRum` instance:
 

Diff for: splunk-otel-android/src/main/java/com/splunk/rum/Config.java

@@ -51,6 +51,8 @@ public class Config {
     private final Duration slowRenderingDetectionPollInterval;
     private final boolean diskBufferingEnabled;
     private final int maxUsageMegabytes;
+    private final boolean sessionBasedSamplerEnabled;
+    private final double sessionBasedSamplerRatio;
 
     private Config(Builder builder) {
         this.beaconEndpoint = builder.beaconEndpoint;
@@ -66,6 +68,8 @@ private Config(Builder builder) {
         this.spanFilterExporterDecorator = builder.spanFilterBuilder.build();
         this.diskBufferingEnabled = builder.diskBufferingEnabled;
         this.maxUsageMegabytes = builder.maxUsageMegabytes;
+        this.sessionBasedSamplerEnabled = builder.sessionBasedSamplerEnabled;
+        this.sessionBasedSamplerRatio = builder.sessionBasedSamplerRatio;
     }
 
     private Attributes addDeploymentEnvironment(Builder builder) {
@@ -167,6 +171,20 @@ public int getMaxUsageMegabytes() {
         return maxUsageMegabytes;
     }
 
+    /**
+     * Is session-based sampling of traces enabled or not.
+     */
+    public boolean isSessionBasedSamplerEnabled() {
+        return sessionBasedSamplerEnabled;
+    }
+
+    /**
+     * Get ratio of sessions that get sampled (0.0 - 1.0, where 1 is all sessions).
+     */
+    public double getSessionBasedSamplerRatio() {
+        return sessionBasedSamplerRatio;
+    }
+
     /**
      * Create a new instance of the {@link Builder} class. All default configuration options will be pre-populated.
      */
@@ -214,6 +232,8 @@ public static class Builder {
         private String realm;
         private Duration slowRenderingDetectionPollInterval = DEFAULT_SLOW_RENDERING_DETECTION_POLL_INTERVAL;
         private int maxUsageMegabytes = DEFAULT_MAX_STORAGE_USE_MB;
+        private boolean sessionBasedSamplerEnabled = false;
+        private double sessionBasedSamplerRatio = 1.0;
 
         /**
          * Create a new instance of {@link Config} from the options provided.
@@ -401,5 +421,35 @@ public Builder limitDiskUsageMegabytes(int maxUsageMegabytes) {
             this.maxUsageMegabytes = maxUsageMegabytes;
             return this;
         }
+
+        /**
+         * Enable/disable session-based sampling of traces. Disabled by default.
+         *
+         * @return {@code this}.
+         */
+        public Builder sessionBasedSamplingEnabled(boolean enabled) {
+            this.sessionBasedSamplerEnabled = enabled;
+            return this;
+        }
+
+        /**
+         * Set ratio of sessions that get sampled (0.0 - 1.0, where 1 is all sessions). Default is
+         * 1.0.
+         *
+         * @return {@code this}.
+         */
+        public Builder enableSessionBasedSampling(double ratio) {
+            if (ratio < 0.0) {
+                Log.e(SplunkRum.LOG_TAG, "invalid sessionBasedSamplingRatio: " + ratio + " must not be negative");
+                return this;
+            } else if (ratio > 1.0) {
+                Log.e(SplunkRum.LOG_TAG, "invalid sessionBasedSamplingRatio: " + ratio + " must not be greater than 1.0");
+                return this;
+            }
+
+            this.sessionBasedSamplerEnabled = true;
+            this.sessionBasedSamplerRatio = ratio;
+            return this;
+        }
     }
 }

Diff for: splunk-otel-android/src/main/java/com/splunk/rum/RumInitializer.java

@@ -33,7 +33,6 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.Locale;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -238,6 +237,10 @@ private SdkTracerProvider buildTracerProvider(
                 .setResource(resource);
         initializationEvents.add(new RumInitializer.InitializationEvent("tracerProviderBuilderInitialized", timingClock.now()));
 
+        if (config.isSessionBasedSamplerEnabled()) {
+            tracerProviderBuilder.setSampler(new SessionIdRatioBasedSampler(config.getSessionBasedSamplerRatio(), sessionId));
+        }
+
         if (config.isDebugEnabled()) {
             tracerProviderBuilder.addSpanProcessor(
                     SimpleSpanProcessor.create(

Diff for: splunk-otel-android/src/main/java/com/splunk/rum/SessionIdRatioBasedSampler.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.rum;
+
+import java.util.List;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+
+/**
+ * Session ID ratio based sampler. Uses {@link Sampler#traceIdRatioBased(double)} sampler
+ * internally, but passes sessionId instead of traceId to the underlying sampler in order to use the
+ * same ratio logic but on sessionId instead. This is valid as {@link SessionId} uses
+ * {@link io.opentelemetry.api.trace.TraceId#fromLongs(long, long)} internally to generate random
+ * session IDs.
+ */
+class SessionIdRatioBasedSampler implements Sampler {
+    private final SessionId sessionId;
+    private final Sampler ratioBasedSampler;
+
+    SessionIdRatioBasedSampler(double ratio, SessionId sessionId) {
+        this.sessionId = sessionId;
+        // SessionId uses the same format as TraceId, so we can reuse trace ID ratio sampler.
+        this.ratioBasedSampler = Sampler.traceIdRatioBased(ratio);
+    }
+
+    @Override
+    public SamplingResult shouldSample(Context parentContext, String traceId, String name, SpanKind spanKind, Attributes attributes, List<LinkData> parentLinks) {
+        // Replace traceId with sessionId
+        return ratioBasedSampler.shouldSample(parentContext, sessionId.getSessionId(), name, spanKind, attributes, parentLinks);
+    }
+
+    @Override
+    public String getDescription() {
+        return String.format("SessionIdRatioBased{traceIdRatioBased:%s}", this.ratioBasedSampler.getDescription());
+    }
+}

Diff for: splunk-otel-android/src/test/java/com/splunk/rum/SessionIdRatioBasedSamplerTest.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.rum;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.IdGenerator;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+
+@RunWith(MockitoJUnitRunner.class)
+public class SessionIdRatioBasedSamplerTest {
+    private static final String HIGH_ID = "00000000000000008fffffffffffffff";
+    private static final String LOW_ID = "00000000000000000000000000000000";
+    private static final IdGenerator idsGenerator = IdGenerator.random();
+
+    private final String traceId = idsGenerator.generateTraceId();
+    private final Context parentContext = Context.root().with(Span.getInvalid());
+    private final List<LinkData> parentLinks = Collections.singletonList(LinkData.create(SpanContext.getInvalid()));
+
+    @Test
+    public void samplerUsesSessionId() {
+        SessionId sessionId = mock(SessionId.class);
+        SessionIdRatioBasedSampler sampler = new SessionIdRatioBasedSampler(0.5, sessionId);
+
+        // Sampler drops if TraceIdRatioBasedSampler would drop this sessionId
+        when(sessionId.getSessionId()).thenReturn(HIGH_ID);
+        assertEquals(shouldSample(sampler), SamplingDecision.DROP);
+
+        // Sampler accepts if TraceIdRatioBasedSampler would accept this sessionId
+        when(sessionId.getSessionId()).thenReturn(LOW_ID);
+        assertEquals(shouldSample(sampler), SamplingDecision.RECORD_AND_SAMPLE);
+    }
+
+    @Test
+    public void zeroRatioDropsAll() {
+        SessionId sessionId = mock(SessionId.class);
+        SessionIdRatioBasedSampler sampler = new SessionIdRatioBasedSampler(0.0, sessionId);
+
+        for (String id : Arrays.asList(HIGH_ID, LOW_ID)) {
+            when(sessionId.getSessionId()).thenReturn(id);
+            assertEquals(shouldSample(sampler), SamplingDecision.DROP);
+        }
+    }
+
+    @Test
+    public void oneRatioAcceptsAll() {
+        SessionId sessionId = mock(SessionId.class);
+        SessionIdRatioBasedSampler sampler = new SessionIdRatioBasedSampler(1.0, sessionId);
+        
+        for (String id : Arrays.asList(HIGH_ID, LOW_ID)) {
+            when(sessionId.getSessionId()).thenReturn(id);
+            assertEquals(shouldSample(sampler), SamplingDecision.RECORD_AND_SAMPLE);
+        }
+    }
+
+    private SamplingDecision shouldSample(Sampler sampler) {
+        return sampler.shouldSample(parentContext, traceId, "name", SpanKind.INTERNAL, Attributes.empty(), parentLinks).getDecision();
+    }
+}