fix: Fix the race condition in decay average

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlStats.java

@@ -15,48 +15,52 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
+import com.google.api.core.ApiClock;
 import com.google.api.gax.batching.FlowController;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Preconditions;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import org.threeten.bp.Instant;
 
 /**
  * Records stats used in dynamic flow control, the decaying average of recorded latencies and the
  * last timestamp when the thresholds in {@link FlowController} are updated.
+ *
+ * <pre>Exponential decaying average = weightedSum / weightedCount, where
+ *   weightedSum(n) = weight(n) * value(n) + weightedSum(n - 1)
+ *   weightedCount(n) = weight(n) + weightedCount(n - 1),
+ * and weight(n) grows exponentially over elapsed time.
  */
 final class DynamicFlowControlStats {
 
-  private static final double DEFAULT_DECAY_CONSTANT = 0.015; // Biased to the past 5 minutes
+  // Biased to the past 5 minutes (300 seconds), e^(-decay_constant * 300) = 0.01, decay_constant ~=
+  // 0.015
+  private static final double DEFAULT_DECAY_CONSTANT = 0.015;
+  // Update start time every 15 minutes so the values won't be infinite
+  private static final long UPDATE_START_TIME_THRESHOLD_SECOND = TimeUnit.MINUTES.toSeconds(15);
 
   private AtomicLong lastAdjustedTimestampMs;
   private DecayingAverage meanLatency;
 
   DynamicFlowControlStats() {
-    this(DEFAULT_DECAY_CONSTANT);
+    this(DEFAULT_DECAY_CONSTANT, null);
   }
 
-  DynamicFlowControlStats(double decayConstant) {
-    this.lastAdjustedTimestampMs = new AtomicLong(0);
-    this.meanLatency = new DecayingAverage(decayConstant);
+  DynamicFlowControlStats(ApiClock clock) {
+    this(DEFAULT_DECAY_CONSTANT, clock);
   }
 
-  void updateLatency(long latency) {
-    updateLatency(latency, System.currentTimeMillis());
+  DynamicFlowControlStats(double decayConstant, ApiClock clock) {
+    this.lastAdjustedTimestampMs = new AtomicLong(0);
+    this.meanLatency = new DecayingAverage(decayConstant, clock);
   }
 
-  @VisibleForTesting
-  void updateLatency(long latency, long timestampMs) {
-    meanLatency.update(latency, timestampMs);
+  void updateLatency(long latency) {
+    meanLatency.update(latency);
   }
 
+  /** Return the mean calculated from the last update, will not decay over time. */
   double getMeanLatency() {
-    return getMeanLatency(System.currentTimeMillis());
-  }
-
-  @VisibleForTesting
-  double getMeanLatency(long timestampMs) {
-    return meanLatency.getMean(timestampMs);
+    return meanLatency.getMean();
   }
 
   public long getLastAdjustedTimestampMs() {
@@ -71,46 +75,58 @@ private class DecayingAverage {
     private double decayConstant;
     private double mean;
     private double weightedCount;
-    private AtomicLong lastUpdateTimeInSecond;
+    private long startTimeSecond;
+    private ApiClock clock;
 
-    DecayingAverage(double decayConstant) {
+    DecayingAverage(double decayConstant, ApiClock clock) {
       this.decayConstant = decayConstant;
       this.mean = 0.0;
       this.weightedCount = 0.0;
-      this.lastUpdateTimeInSecond = new AtomicLong(0);
+      this.clock = clock;
+      this.startTimeSecond =
+          clock == null
+              ? Instant.now().getEpochSecond()
+              : TimeUnit.MILLISECONDS.toSeconds(clock.millisTime());
     }
 
-    synchronized void update(long value, long timestampMs) {
-      long now = TimeUnit.MILLISECONDS.toSeconds(timestampMs);
-      Preconditions.checkArgument(
-          now >= lastUpdateTimeInSecond.get(), "can't update an event in the past");
-      if (lastUpdateTimeInSecond.get() == 0) {
-        lastUpdateTimeInSecond.set(now);
-        mean = value;
-        weightedCount = 1;
-      } else {
-        long prev = lastUpdateTimeInSecond.getAndSet(now);
-        long elapsed = now - prev;
-        double alpha = getAlpha(elapsed);
-        // Exponential moving average = weightedSum / weightedCount, where
-        // weightedSum(n) = value + alpha * weightedSum(n - 1)
-        // weightedCount(n) = 1 + alpha * weightedCount(n - 1)
-        // Using weighted count in case the sum overflows
-        mean =
-            mean * ((weightedCount * alpha) / (weightedCount * alpha + 1))
-                + value / (weightedCount * alpha + 1);
-        weightedCount = weightedCount * alpha + 1;
-      }
+    synchronized void update(long value) {
+      // Decay mean and weightedCount if now - startTime > threshold, so weight won't be infinite
+      long now = getCurrentTimeInSecond();
+      double decay = getDecay(now);
+      mean /= decay;
+      weightedCount /= decay;
+
+      long elapsed = now - startTimeSecond;
+      double weight = getWeight(elapsed);
+      // Using weighted count in case the sum overflows.
+      mean =
+          mean * (weightedCount / (weightedCount + weight))
+              + weight * value / (weightedCount + weight);
+      weightedCount += weight;
+    }
+
+    double getMean() {
+      return mean;
     }
 
-    double getMean(long timestampMs) {
-      long timestampSecs = TimeUnit.MILLISECONDS.toSeconds(timestampMs);
-      long elapsed = timestampSecs - lastUpdateTimeInSecond.get();
-      return mean * getAlpha(Math.max(0, elapsed));
+    private long getCurrentTimeInSecond() {
+      return clock == null
+          ? Instant.now().getEpochSecond()
+          : TimeUnit.MILLISECONDS.toSeconds(clock.millisTime());
     }
 
-    private double getAlpha(long elapsedSecond) {
-      return Math.exp(-decayConstant * elapsedSecond);
+    private double getWeight(long elapsedSecond) {
+      return Math.exp(decayConstant * elapsedSecond);
+    }
+
+    private synchronized double getDecay(long now) {
+      long elapsed = now - startTimeSecond;
+      if (elapsed > UPDATE_START_TIME_THRESHOLD_SECOND) {
+        double decay = getWeight(elapsed);
+        startTimeSecond = now;
+        return decay;
+      }
+      return 1;
     }
   }
 }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlStatsTest.java

@@ -27,40 +27,39 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.threeten.bp.Instant;
 
 @RunWith(JUnit4.class)
 public class DynamicFlowControlStatsTest {
 
   @Test
   public void testUpdate() {
-    DynamicFlowControlStats stats = new DynamicFlowControlStats();
-    long now = System.currentTimeMillis();
+    FakeApiClock clock = new FakeApiClock(Instant.now().toEpochMilli());
 
-    stats.updateLatency(10, now);
-    assertThat(stats.getMeanLatency(now)).isEqualTo(10);
+    DynamicFlowControlStats stats = new DynamicFlowControlStats(clock);
 
-    stats.updateLatency(10, now);
-    stats.updateLatency(10, now);
-    assertThat(stats.getMeanLatency(now)).isEqualTo(10);
+    stats.updateLatency(10);
+    assertThat(stats.getMeanLatency()).isEqualTo(10);
+    stats.updateLatency(10);
+    stats.updateLatency(10);
+    assertThat(stats.getMeanLatency()).isEqualTo(10);
 
     // In five minutes the previous latency should be decayed to under 1. And the new average should
     // be very close to 20
-    long fiveMinutesLater = now + TimeUnit.MINUTES.toMillis(5);
-    assertThat(stats.getMeanLatency(fiveMinutesLater)).isLessThan(1);
-    stats.updateLatency(20, fiveMinutesLater);
-    assertThat(stats.getMeanLatency(fiveMinutesLater)).isGreaterThan(19);
-    assertThat(stats.getMeanLatency(fiveMinutesLater)).isLessThan(20);
+    clock.tick(TimeUnit.MINUTES.toMillis(5));
+    stats.updateLatency(20);
+    assertThat(stats.getMeanLatency()).isGreaterThan(19);
+    assertThat(stats.getMeanLatency()).isLessThan(20);
 
-    long aDayLater = now + TimeUnit.HOURS.toMillis(24);
-    assertThat(stats.getMeanLatency(aDayLater)).isZero();
+    // After a day
+    clock.tick(TimeUnit.HOURS.toMillis(24));
 
-    long timestamp = aDayLater;
     for (int i = 0; i < 10; i++) {
-      timestamp += TimeUnit.SECONDS.toMillis(i);
-      stats.updateLatency(i, timestamp);
+      clock.tick(TimeUnit.SECONDS.toMillis(i));
+      stats.updateLatency(i);
     }
-    assertThat(stats.getMeanLatency(timestamp)).isGreaterThan(4.5);
-    assertThat(stats.getMeanLatency(timestamp)).isLessThan(6);
+    assertThat(stats.getMeanLatency()).isGreaterThan(4.5);
+    assertThat(stats.getMeanLatency()).isLessThan(6);
   }
 
   @Test(timeout = 1000)

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/FakeApiClock.java

@@ -0,0 +1,43 @@
+/*
+ * 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
+ *
+ *     https://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.cloud.bigtable.data.v2.stub;
+
+import com.google.api.core.ApiClock;
+import java.util.concurrent.TimeUnit;
+
+public class FakeApiClock implements ApiClock {
+
+  private long millisTime;
+
+  public FakeApiClock(long millisTime) {
+    this.millisTime = millisTime;
+  }
+
+  @Override
+  public long nanoTime() {
+    return TimeUnit.MILLISECONDS.toNanos(millisTime);
+  }
+
+  @Override
+  public long millisTime() {
+    return millisTime;
+  }
+
+  public void tick(long millis) {
+    this.millisTime += millis;
+  }
+}