Rate limiting should be ineffective when RateLimitInfo is not present

Built a new ConditionalRateLimiter that can be disabled but keep the rate. Also moved setRate() into the ConditionalRateLimiter. Update and disable is throttled based on period from RateLimitInfo.

Author: kongweihan

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

@@ -31,15 +31,16 @@
 import com.google.common.base.Stopwatch;
 import com.google.common.util.concurrent.RateLimiter;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nonnull;
 import org.threeten.bp.Duration;
 import org.threeten.bp.Instant;
 
 class RateLimitingServerStreamingCallable
  extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
+
  private static final Logger logger =
  Logger.getLogger(RateLimitingServerStreamingCallable.class.getName());
 
@@ -64,16 +65,17 @@ class RateLimitingServerStreamingCallable
  // as the server side cap
  private static final double MAX_FACTOR = 1.3;
 
- private final RateLimiter limiter;
+ private final ConditionalRateLimiter limiter;
 
- private final AtomicReference<Instant> lastQpsChangeTime = new AtomicReference<>(Instant.now());
  private final ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable;
 
  RateLimitingServerStreamingCallable(
  @Nonnull ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable) {
- this.limiter = RateLimiter.create(DEFAULT_QPS);
+ this.limiter = new ConditionalRateLimiter(DEFAULT_QPS);
  this.innerCallable = Preconditions.checkNotNull(innerCallable, "Inner callable must be set");
- logger.info("Rate limiting is enabled with initial QPS of " + limiter.getRate());
+ logger.info(
+ "BatchWriteFlowControl: Rate limiting callable is initiated with QPS of "
+ + limiter.getRate());
  }
 
  @Override
@@ -88,44 +90,129 @@ public void call(
  ((BigtableTracer) context.getTracer())
  .batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
  }
- RateLimitingResponseObserver innerObserver =
- new RateLimitingResponseObserver(limiter, lastQpsChangeTime, responseObserver);
+ RateLimitingResponseObserver innerObserver = new RateLimitingResponseObserver(responseObserver);
  innerCallable.call(request, innerObserver, context);
  }
 
+ static class ConditionalRateLimiter {
+
+ private AtomicBoolean enabled = new AtomicBoolean(true);
+
+ private final RateLimiter limiter;
+
+ // This is the next time allowed to change QPS or disable rate limiting.
+ private final AtomicReference<Instant> nextQpsChangeTime = new AtomicReference<>(Instant.now());
+
+ public ConditionalRateLimiter(long defaultQps) {
+ limiter = RateLimiter.create(defaultQps);
+ }
+
+ public double acquire() {
+ if (enabled.get()) {
+ return limiter.acquire();
+ } else {
+ return 0;
+ }
+ }
+
+ // Enable rate limiting immediately or disable after the QPS update period. Otherwise, no-op.
+ public void trySetEnabled(boolean enabled) {
+ if (enabled) {
+ // Always enable immediately.
+ boolean wasEnabled = this.enabled.getAndSet(true);
+ if (!wasEnabled) {
+ logger.fine("BatchWriteFlowControl: rate limiter is enabled.");
+ }
+ return;
+ }
+ // Only disable after the QPS update period.
+ Instant nextTime = nextQpsChangeTime.get();
+ Instant now = Instant.now();
+ if (now.isAfter(nextTime)) {
+ boolean wasEnabled = this.enabled.getAndSet(false);
+ if (wasEnabled) {
+ logger.fine("BatchWriteFlowControl: rate limiter is disabled.");
+ }
+ }
+ }
+
+ public boolean isEnabled() {
+ return this.enabled.get();
+ }
+
+ public double getRate() {
+ return limiter.getRate();
+ }
+
+ // Update the rate after the QPS update period. Otherwise, no-op.
+ public void trySetRate(double rate, Duration period) {
+ Instant nextTime = nextQpsChangeTime.get();
+ Instant now = Instant.now();
+
+ if (now.isAfter(nextTime)
+ && nextQpsChangeTime.compareAndSet(nextTime, now.plusSeconds(period.getSeconds()))) {
+ logger.fine(
+ "BatchWriteFlowControl: updated rate from "
+ + limiter.getRate()
+ + " to "
+ + rate
+ + " with period "
+ + period.getSeconds()
+ + " seconds.");
+ limiter.setRate(rate);
+ }
+ }
+
+ @VisibleForTesting
+ void setEnabled(boolean enabled) {
+ this.enabled.set(enabled);
+ }
+ }
+
  class RateLimitingResponseObserver extends SafeResponseObserver<MutateRowsResponse> {
- private final ResponseObserver<MutateRowsResponse> outerObserver;
- private final RateLimiter rateLimiter;
 
- private final AtomicReference<Instant> lastQpsChangeTime;
+ private final ResponseObserver<MutateRowsResponse> outerObserver;
 
- RateLimitingResponseObserver(
- RateLimiter rateLimiter,
- AtomicReference<Instant> lastQpsChangeTime,
- ResponseObserver<MutateRowsResponse> observer) {
+ RateLimitingResponseObserver(ResponseObserver<MutateRowsResponse> observer) {
  super(observer);
  this.outerObserver = observer;
- this.rateLimiter = rateLimiter;
- this.lastQpsChangeTime = lastQpsChangeTime;
  }
 
  @Override
  protected void onStartImpl(StreamController controller) {
  outerObserver.onStart(controller);
  }
 
+ private boolean isRateLimitInfoAbsentOrInvalid(MutateRowsResponse response) {
+ // RateLimitInfo is an optional field. However, proto3 sub-message field always
+ // have presence even thought it's marked as "optional". Check the factor and
+ // period to make sure they're not 0.
+ if (!response.hasRateLimitInfo()) {
+ logger.fine("BatchWriteFlowControl: response carries no RateLimitInfo");
+ return true;
+ } else if (response.getRateLimitInfo().getFactor() == 0
+ || response.getRateLimitInfo().getPeriod().getSeconds() == 0) {
+ logger.fine(
+ "BatchWriteFlowControl: response carries invalid RateLimitInfo="
+ + response.getRateLimitInfo());
+ return true;
+ }
+ logger.fine(
+ "BatchWriteFlowControl: response carries valid RateLimitInfo="
+ + response.getRateLimitInfo());
+ return false;
+ }
+
  @Override
  protected void onResponseImpl(MutateRowsResponse response) {
- if (response.hasRateLimitInfo()) {
+ if (isRateLimitInfoAbsentOrInvalid(response)) {
+ limiter.trySetEnabled(false);
+ } else {
+ limiter.trySetEnabled(true);
  RateLimitInfo info = response.getRateLimitInfo();
- // RateLimitInfo is an optional field. However, proto3 sub-message field always
- // have presence even thought it's marked as "optional". Check the factor and
- // period to make sure they're not 0.
- if (info.getFactor() != 0 && info.getPeriod().getSeconds() != 0) {
- updateQps(
- info.getFactor(),
- Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
- }
+ updateQps(
+ info.getFactor(),
+ Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
  }
  outerObserver.onResponse(response);
  }
@@ -148,28 +235,30 @@ protected void onCompleteImpl() {
  }
 
  private void updateQps(double factor, Duration period) {
- Instant lastTime = lastQpsChangeTime.get();
- Instant now = Instant.now();
-
- if (now.minus(period).isAfter(lastTime) && lastQpsChangeTime.compareAndSet(lastTime, now)) {
- double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR);
- double currentRate = limiter.getRate();
- limiter.setRate(Math.min(Math.max(currentRate * cappedFactor, MIN_QPS), MAX_QPS));
- logger.log(
- Level.FINE,
- "Updated QPS from {0} to {1}, server returned factor is {2}, capped factor is {3}",
- new Object[] {currentRate, limiter.getRate(), factor, cappedFactor});
- }
+ double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR);
+ double currentRate = limiter.getRate();
+ double cappedRate = Math.min(Math.max(currentRate * cappedFactor, MIN_QPS), MAX_QPS);
+ limiter.trySetRate(cappedRate, period);
  }
  }
 
  @VisibleForTesting
- AtomicReference<Instant> getLastQpsChangeTime() {
- return lastQpsChangeTime;
+ AtomicReference<Instant> getNextQpsChangeTime() {
+ return limiter.nextQpsChangeTime;
  }
 
  @VisibleForTesting
  double getCurrentRate() {
  return limiter.getRate();
  }
+
+ @VisibleForTesting
+ boolean getLimiterEnabled() {
+ return limiter.isEnabled();
+ }
+
+ @VisibleForTesting
+ void setLimiterEnabled(boolean enabled) {
+ limiter.setEnabled(enabled);
+ }
 }