Fix fireperf RateLimiter

Author: jfcong

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

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.perf.metrics.resource.ResourceType.NETWORK;
 import static com.google.firebase.perf.metrics.resource.ResourceType.TRACE;
+import static java.util.concurrent.TimeUnit.SECONDS;
 
 import android.content.Context;
 import androidx.annotation.NonNull;
@@ -25,6 +26,7 @@
 import com.google.firebase.perf.metrics.resource.ResourceType;
 import com.google.firebase.perf.util.Clock;
 import com.google.firebase.perf.util.Constants;
+import com.google.firebase.perf.util.Rate;
 import com.google.firebase.perf.util.Timer;
 import com.google.firebase.perf.util.Utils;
 import com.google.firebase.perf.v1.NetworkRequestMetric;
@@ -34,7 +36,6 @@
 import com.google.firebase.perf.v1.TraceMetric;
 import java.util.List;
 import java.util.Random;
-import java.util.concurrent.TimeUnit;
 
 /**
  * Implement the Token Bucket rate limiting algorithm. The token bucket initially holds "capacity"
@@ -60,10 +61,10 @@ final class RateLimiter {
    * Construct a token bucket rate limiter.
    *
    * @param appContext the application's context.
-   * @param rate number of token generated per minute.
+   * @param rate the Rate object representing the number of tokens generated per specified time unit
    * @param capacity token bucket capacity
    */
-  public RateLimiter(@NonNull Context appContext, final double rate, final long capacity) {
+  public RateLimiter(@NonNull Context appContext, final Rate rate, final long capacity) {
     this(rate, capacity, new Clock(), getSamplingBucketId(), ConfigResolver.getInstance());
     this.isLogcatEnabled = Utils.isDebugLoggingEnabled(appContext);
   }
@@ -75,7 +76,7 @@ static float getSamplingBucketId() {
   }
 
   RateLimiter(
-      final double rate,
+      final Rate rate,
       final long capacity,
       final Clock clock,
       float samplingBucketId,
@@ -211,29 +212,29 @@ boolean getIsDeviceAllowedToSendNetworkEvents() {
   static class RateLimiterImpl {
 
     private static final AndroidLogger logger = AndroidLogger.getInstance();
-    private static final long MICROS_IN_A_SECOND = TimeUnit.SECONDS.toMicros(1);
+    private static final long MICROS_IN_A_SECOND = SECONDS.toMicros(1);
 
     private final Clock clock;
     private final boolean isLogcatEnabled;
 
     // Last time a token is consumed.
-    private Timer lastTimeTokenConsumed;
+    private Timer lastTimeTokenReplenished;
 
-    // Number of new tokens generated per second.
-    private double rate;
+    // The Rate object representing the number of tokens generated per specified time unit
+    private Rate rate;
     // Token bucket capacity, also the initial number of tokens in the bucket.
     private long capacity;
     // Number of tokens in the bucket.
     private long tokenCount;
 
-    private double foregroundRate;
-    private double backgroundRate;
+    private Rate foregroundRate;
+    private Rate backgroundRate;
 
     private long foregroundCapacity;
     private long backgroundCapacity;
 
     RateLimiterImpl(
-        final double rate,
+        final Rate rate,
         final long capacity,
         final Clock clock,
         ConfigResolver configResolver,
@@ -243,7 +244,7 @@ static class RateLimiterImpl {
       this.capacity = capacity;
       this.rate = rate;
       tokenCount = capacity;
-      lastTimeTokenConsumed = this.clock.getTime();
+      lastTimeTokenReplenished = this.clock.getTime();
       setRateByReadingRemoteConfigValues(configResolver, type, isLogcatEnabled);
       this.isLogcatEnabled = isLogcatEnabled;
     }
@@ -260,11 +261,20 @@ synchronized boolean check(@NonNull PerfMetric metric) {
       Timer now = clock.getTime();
       long newTokens =
           Math.max(
-              0, (long) (lastTimeTokenConsumed.getDurationMicros(now) * rate / MICROS_IN_A_SECOND));
+              0,
+              (long)
+                  (lastTimeTokenReplenished.getDurationMicros(now)
+                      * rate.getTokenPerSeconds()
+                      / MICROS_IN_A_SECOND));
       tokenCount = Math.min(tokenCount + newTokens, capacity);
+      if (newTokens > 0) {
+        lastTimeTokenReplenished =
+            new Timer(
+                lastTimeTokenReplenished.getMicros()
+                    + (long) (newTokens * MICROS_IN_A_SECOND / rate.getTokenPerSeconds()));
+      }
       if (tokenCount > 0) {
         tokenCount--;
-        lastTimeTokenConsumed = now;
         return true;
       }
       if (isLogcatEnabled) {
@@ -297,7 +307,7 @@ private void setRateByReadingRemoteConfigValues(
       long fLimitTime = getFlimitSec(configResolver, type);
       long fLimitEvents = getFlimitEvents(configResolver, type);
 
-      foregroundRate = ((double) fLimitEvents) / fLimitTime;
+      foregroundRate = new Rate(fLimitEvents, fLimitTime, SECONDS);
       foregroundCapacity = fLimitEvents;
       if (isLogcatEnabled) {
         logger.debug(
@@ -309,7 +319,7 @@ private void setRateByReadingRemoteConfigValues(
       long bLimitTime = getBlimitSec(configResolver, type);
       long bLimitEvents = getBlimitEvents(configResolver, type);
 
-      backgroundRate = ((double) bLimitEvents) / bLimitTime;
+      backgroundRate = new Rate(bLimitEvents, bLimitTime, SECONDS);
       backgroundCapacity = bLimitEvents;
       if (isLogcatEnabled) {
         logger.debug(
@@ -350,7 +360,7 @@ private static long getBlimitEvents(
     }
 
     @VisibleForTesting
-    double getForegroundRate() {
+    Rate getForegroundRate() {
       return foregroundRate;
     }
 
@@ -360,7 +370,7 @@ long getForegroundCapacity() {
     }
 
     @VisibleForTesting
-    double getBackgroundRate() {
+    Rate getBackgroundRate() {
       return backgroundRate;
     }
 
@@ -370,12 +380,12 @@ long getBackgroundCapacity() {
     }
 
     @VisibleForTesting
-    double getRate() {
+    Rate getRate() {
       return rate;
     }
 
     @VisibleForTesting
-    void setRate(double newRate) {
+    void setRate(Rate newRate) {
       rate = newRate;
     }
   }

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

@@ -15,6 +15,7 @@
 package com.google.firebase.perf.transport;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.MINUTES;
 
 import android.content.Context;
 import android.content.pm.PackageInfo;
@@ -39,6 +40,7 @@
 import com.google.firebase.perf.session.SessionManager;
 import com.google.firebase.perf.util.Constants;
 import com.google.firebase.perf.util.Constants.CounterNames;
+import com.google.firebase.perf.util.Rate;
 import com.google.firebase.perf.v1.AndroidApplicationInfo;
 import com.google.firebase.perf.v1.ApplicationInfo;
 import com.google.firebase.perf.v1.ApplicationProcessState;
@@ -202,7 +204,9 @@ public void initialize(
   private void syncInit() {
     appContext = firebaseApp.getApplicationContext();
     configResolver = ConfigResolver.getInstance();
-    rateLimiter = new RateLimiter(appContext, Constants.RATE_PER_MINUTE, Constants.BURST_CAPACITY);
+    rateLimiter =
+        new RateLimiter(
+            appContext, new Rate(Constants.RATE_PER_MINUTE, 1, MINUTES), Constants.BURST_CAPACITY);
     appStateMonitor = AppStateMonitor.getInstance();
     flgTransport =
         new FlgTransport(flgTransportFactoryProvider, configResolver.getAndCacheLogSourceName());

firebase-perf/src/main/java/com/google/firebase/perf/util/Rate.java

@@ -0,0 +1,64 @@
+// Copyright 2021 Google LLC
+//
+// 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.google.firebase.perf.util;
+
+import java.util.concurrent.TimeUnit;
+
+/** A Rate object representing the number of tokens per total specified time unit. */
+public class Rate {
+
+  private static final long NANO_IN_SECOND = 1000 * 1000 * 1000;
+  private static final long MICRO_IN_SECOND = 1000 * 1000;
+  private static final long MILLI_IN_SECOND = 1000;
+
+  private long numTokensPerTotalTimeUnit;
+  private long numTimeUnits;
+  private TimeUnit timeUnit;
+
+  /**
+   * Constructs a Rate object.
+   *
+   * <p>For example, a rate of 3 tokens per minute can be represented by new Rate(3, 1, MINUTES) or
+   * new Rate(3, 60, SECONDS).
+   *
+   * @param numTokensPerTotalTimeUnit The number of tokens to be issued within the specified time
+   * @param numTimeUnits The number of specified time unit
+   * @param timeUnit The specified time unit
+   */
+  public Rate(long numTokensPerTotalTimeUnit, long numTimeUnits, TimeUnit timeUnit) {
+    assert numTimeUnits > 0;
+    this.numTokensPerTotalTimeUnit = numTokensPerTotalTimeUnit;
+    this.numTimeUnits = numTimeUnits;
+    this.timeUnit = timeUnit;
+  }
+
+  /**
+   * Converts the rate to tokens per second.
+   *
+   * @return rate in tokens per second
+   */
+  public double getTokenPerSeconds() {
+    switch (timeUnit) {
+      case NANOSECONDS:
+        return ((double) numTokensPerTotalTimeUnit / numTimeUnits) * NANO_IN_SECOND;
+      case MICROSECONDS:
+        return ((double) numTokensPerTotalTimeUnit / numTimeUnits) * MICRO_IN_SECOND;
+      case MILLISECONDS:
+        return ((double) numTokensPerTotalTimeUnit / numTimeUnits) * MILLI_IN_SECOND;
+      default:
+        return (double) numTokensPerTotalTimeUnit / timeUnit.toSeconds(numTimeUnits);
+    }
+  }
+}