Fireperf fragments: sampling

firebase-perf/src/main/java/com/google/firebase/perf/transport/RateLimiter.java

@@ -49,6 +49,8 @@ final class RateLimiter {
   /** The app's bucket ID for sampling, a number in [0.0f, 1.0f). */
   private final float samplingBucketId;
 
+  private final float fragmentBucketId;
+
   private RateLimiterImpl traceLimiter = null;
   private RateLimiterImpl networkLimiter = null;
 
@@ -63,7 +65,13 @@ final class RateLimiter {
    * @param capacity token bucket capacity
    */
   public RateLimiter(@NonNull Context appContext, final Rate rate, final long capacity) {
-    this(rate, capacity, new Clock(), getSamplingBucketId(), ConfigResolver.getInstance());
+    this(
+        rate,
+        capacity,
+        new Clock(),
+        getSamplingBucketId(),
+        getSamplingBucketId(),
+        ConfigResolver.getInstance());
     this.isLogcatEnabled = Utils.isDebugLoggingEnabled(appContext);
   }
 
@@ -78,11 +86,16 @@ static float getSamplingBucketId() {
       final long capacity,
       final Clock clock,
       float samplingBucketId,
+      float fragmentBucketId,
       ConfigResolver configResolver) {
     Utils.checkArgument(
         0.0f <= samplingBucketId && samplingBucketId < 1.0f,
         "Sampling bucket ID should be in range [0.0f, 1.0f).");
+    Utils.checkArgument(
+        0.0f <= fragmentBucketId && fragmentBucketId < 1.0f,
+        "Fragment sampling bucket ID should be in range [0.0f, 1.0f).");
     this.samplingBucketId = samplingBucketId;
+    this.fragmentBucketId = fragmentBucketId;
     this.configResolver = configResolver;
 
     traceLimiter =
@@ -104,6 +117,22 @@ private boolean isDeviceAllowedToSendNetworkEvents() {
     return samplingBucketId < validNetworkSamplingBucketIdThreshold;
   }
 
+  /**
+   * Returns whether device is allowed to send Fragment screen trace events based on Fragment screen
+   * trace sampling rate.
+   */
+  private boolean isDeviceAllowedToSendFragmentScreenTraces() {
+    float validFragmentSamplingBucketIdThreshold = configResolver.getFragmentSamplingRate();
+    return fragmentBucketId < validFragmentSamplingBucketIdThreshold;
+  }
+
+  /** Identifies if the {@link PerfMetric} is a Fragment screen trace */
+  protected boolean isFragmentScreenTrace(PerfMetric metric) {
+    return metric.hasTraceMetric()
+        && metric.getTraceMetric().getName().startsWith(Constants.SCREEN_TRACE_PREFIX)
+        && metric.getTraceMetric().containsCustomAttributes(Constants.ACTIVITY_ATTRIBUTE_KEY);
+  }
+
   /**
    * Check if the {@link PerfMetric} should be rate limited.
    *
@@ -140,6 +169,12 @@ boolean isEventSampled(PerfMetric metric) {
       return false;
     }
 
+    if (isFragmentScreenTrace(metric)
+        && !(isDeviceAllowedToSendFragmentScreenTraces()
+            || hasVerboseSessions(metric.getTraceMetric().getPerfSessionsList()))) {
+      return false;
+    }
+
     if (metric.hasNetworkRequestMetric()
         && !(isDeviceAllowedToSendNetworkEvents()
             || hasVerboseSessions(metric.getNetworkRequestMetric().getPerfSessionsList()))) {
@@ -207,6 +242,11 @@ boolean getIsDeviceAllowedToSendNetworkEvents() {
     return isDeviceAllowedToSendNetworkEvents();
   }
 
+  @VisibleForTesting
+  boolean getIsDeviceAllowedToSendFragmentScreenTraces() {
+    return isDeviceAllowedToSendFragmentScreenTraces();
+  }
+
   /** The implementation of Token Bucket rate limiter. */
   static class RateLimiterImpl {