feat: add rate limiting callable (#1729)

* feat: add rate limiting callable

* address comments

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

@@ -0,0 +1,174 @@
+/*
+ * Copyright 2023 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.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.DeadlineExceededException;
+import com.google.api.gax.rpc.ResourceExhaustedException;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.StreamController;
+import com.google.api.gax.rpc.UnavailableException;
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.bigtable.v2.MutateRowsResponse;
+import com.google.bigtable.v2.RateLimitInfo;
+import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
+import com.google.common.util.concurrent.RateLimiter;
+import java.util.concurrent.TimeUnit;
+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());
+
+  // When the mutation size is large, starting with a higher QPS will make
+  // the dataflow job fail very quickly. Start with lower QPS and increase
+  // the QPS gradually if the server doesn't push back
+  private static final long DEFAULT_QPS = 10;
+
+  // Default interval before changing the QPS on error responses
+  private static final Duration DEFAULT_PERIOD = Duration.ofSeconds(10);
+
+  // Minimum QPS to make sure the job is not stuck
+  private static final double MIN_QPS = 0.1;
+  private static final double MAX_QPS = 100_000;
+
+  // QPS can be lowered to at most MIN_FACTOR * currentQps. When server returned
+  // an error, use MIN_FACTOR to calculate the new QPS. This is the same as
+  // the server side cap.
+  @VisibleForTesting static final double MIN_FACTOR = 0.7;
+
+  // QPS can be increased to at most MAX_FACTOR * currentQps. This is the same
+  // as the server side cap
+  private static final double MAX_FACTOR = 1.3;
+
+  private final RateLimiter 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.innerCallable = Preconditions.checkNotNull(innerCallable, "Inner callable must be set");
+    logger.info("Rate limiting is enabled with initial QPS of " + limiter.getRate());
+  }
+
+  @Override
+  public void call(
+      MutateRowsRequest request,
+      ResponseObserver<MutateRowsResponse> responseObserver,
+      ApiCallContext context) {
+    Stopwatch stopwatch = Stopwatch.createStarted();
+    limiter.acquire();
+    stopwatch.stop();
+    if (context.getTracer() instanceof BigtableTracer) {
+      ((BigtableTracer) context.getTracer())
+          .batchRequestThrottled(stopwatch.elapsed(TimeUnit.MILLISECONDS));
+    }
+    RateLimitingResponseObserver innerObserver =
+        new RateLimitingResponseObserver(limiter, lastQpsChangeTime, responseObserver);
+    innerCallable.call(request, innerObserver, context);
+  }
+
+  class RateLimitingResponseObserver extends SafeResponseObserver<MutateRowsResponse> {
+    private final ResponseObserver<MutateRowsResponse> outerObserver;
+    private final RateLimiter rateLimiter;
+
+    private final AtomicReference<Instant> lastQpsChangeTime;
+
+    RateLimitingResponseObserver(
+        RateLimiter rateLimiter,
+        AtomicReference<Instant> lastQpsChangeTime,
+        ResponseObserver<MutateRowsResponse> observer) {
+      super(observer);
+      this.outerObserver = observer;
+      this.rateLimiter = rateLimiter;
+      this.lastQpsChangeTime = lastQpsChangeTime;
+    }
+
+    @Override
+    protected void onStartImpl(StreamController controller) {
+      outerObserver.onStart(controller);
+    }
+
+    @Override
+    protected void onResponseImpl(MutateRowsResponse response) {
+      if (response.hasRateLimitInfo()) {
+        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())));
+        }
+      }
+    }
+
+    @Override
+    protected void onErrorImpl(Throwable t) {
+      // When server returns DEADLINE_EXCEEDED, UNAVAILABLE or RESOURCE_EXHAUSTED,
+      // assume cbt server is overloaded
+      if (t instanceof DeadlineExceededException
+          || t instanceof UnavailableException
+          || t instanceof ResourceExhaustedException) {
+        updateQps(MIN_FACTOR, DEFAULT_PERIOD);
+      }
+      outerObserver.onError(t);
+    }
+
+    @Override
+    protected void onCompleteImpl() {
+      outerObserver.onComplete();
+    }
+
+    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});
+      }
+    }
+  }
+
+  @VisibleForTesting
+  AtomicReference<Instant> getLastQpsChangeTime() {
+    return lastQpsChangeTime;
+  }
+
+  @VisibleForTesting
+  double getCurrentRate() {
+    return limiter.getRate();
+  }
+}

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

@@ -0,0 +1,187 @@
+/*
+ * Copyright 2023 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 static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.DeadlineExceededException;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.StatusCode;
+import com.google.api.gax.rpc.StreamController;
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.bigtable.v2.MutateRowsResponse;
+import com.google.bigtable.v2.RateLimitInfo;
+import com.google.cloud.bigtable.gaxx.testing.FakeStatusCode;
+import com.google.protobuf.Duration;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+import org.threeten.bp.Instant;
+
+@RunWith(JUnit4.class)
+public class RateLimitingCallableTest {
+
+  private final MutateRowsRequest request =
+      MutateRowsRequest.newBuilder().getDefaultInstanceForType();
+  private final ResponseObserver<MutateRowsResponse> responseObserver =
+      Mockito.mock(ResponseObserver.class);
+  private final ApiCallContext context = GrpcCallContext.createDefault();
+  private MockCallable innerCallable;
+  RateLimitingServerStreamingCallable callableToTest;
+
+  @Before
+  public void setup() throws Exception {
+    innerCallable = new MockCallable();
+    callableToTest = new RateLimitingServerStreamingCallable(innerCallable);
+  }
+
+  @Test
+  public void testWithRateLimitInfo() throws Exception {
+    callableToTest.call(request, responseObserver, context);
+
+    Instant earlier = Instant.now().minus(org.threeten.bp.Duration.ofHours(1));
+
+    // make sure QPS will be updated
+    callableToTest.getLastQpsChangeTime().set(earlier);
+    double oldQps = callableToTest.getCurrentRate();
+
+    double factor = 0.8;
+
+    RateLimitInfo info =
+        RateLimitInfo.newBuilder()
+            .setFactor(factor)
+            .setPeriod(Duration.newBuilder().setSeconds(10).build())
+            .build();
+
+    MutateRowsResponse response = MutateRowsResponse.newBuilder().setRateLimitInfo(info).build();
+
+    innerCallable.getObserver().onResponse(response);
+
+    // Give the thread sometime to update the QPS
+    Thread.sleep(100);
+    double newQps = callableToTest.getCurrentRate();
+
+    assertThat(newQps).isWithin(0.1).of(oldQps * factor);
+
+    innerCallable.getObserver().onComplete();
+  }
+
+  @Test
+  public void testNoUpdateWithinPeriod() throws Exception {
+    callableToTest.call(request, responseObserver, context);
+
+    Instant now = Instant.now();
+    // make sure QPS will not be updated
+    callableToTest.getLastQpsChangeTime().set(now);
+    double oldQps = callableToTest.getCurrentRate();
+
+    double factor = 0.3;
+
+    RateLimitInfo info =
+        RateLimitInfo.newBuilder()
+            .setFactor(factor)
+            .setPeriod(Duration.newBuilder().setSeconds(600).build())
+            .build();
+
+    MutateRowsResponse response = MutateRowsResponse.newBuilder().setRateLimitInfo(info).build();
+
+    innerCallable.getObserver().onResponse(response);
+
+    // Give the thread sometime to update the QPS
+    Thread.sleep(100);
+    double newQps = callableToTest.getCurrentRate();
+
+    assertThat(newQps).isEqualTo(oldQps);
+
+    innerCallable.getObserver().onComplete();
+  }
+
+  @Test
+  public void testErrorInfoLowerQPS() throws Exception {
+    callableToTest.call(request, responseObserver, context);
+
+    Instant earlier = Instant.now().minus(org.threeten.bp.Duration.ofHours(1));
+
+    // make sure QPS will be updated
+    callableToTest.getLastQpsChangeTime().set(earlier);
+    double oldQps = callableToTest.getCurrentRate();
+
+    innerCallable
+        .getObserver()
+        .onError(
+            new DeadlineExceededException(
+                new Throwable(), new FakeStatusCode(StatusCode.Code.DEADLINE_EXCEEDED), false));
+
+    // Give the thread sometime to update the QPS
+    Thread.sleep(100);
+    double newQps = callableToTest.getCurrentRate();
+
+    assertThat(newQps).isWithin(0.1).of(oldQps * RateLimitingServerStreamingCallable.MIN_FACTOR);
+  }
+
+  private static class MockResponseObserver implements ResponseObserver<MutateRowsResponse> {
+
+    private ResponseObserver<MutateRowsResponse> observer;
+
+    MockResponseObserver(ResponseObserver<MutateRowsResponse> responseObserver) {
+      this.observer = responseObserver;
+    }
+
+    @Override
+    public void onStart(StreamController streamController) {
+      observer.onStart(streamController);
+    }
+
+    @Override
+    public void onResponse(MutateRowsResponse o) {
+      observer.onResponse(o);
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+      observer.onError(throwable);
+    }
+
+    @Override
+    public void onComplete() {
+      observer.onComplete();
+    }
+  }
+
+  private static class MockCallable
+      extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
+
+    private ResponseObserver<MutateRowsResponse> observer;
+
+    @Override
+    public void call(
+        MutateRowsRequest mutateRowsRequest,
+        ResponseObserver<MutateRowsResponse> responseObserver,
+        ApiCallContext apiCallContext) {
+      observer = new MockResponseObserver(responseObserver);
+    }
+
+    ResponseObserver<MutateRowsResponse> getObserver() {
+      return observer;
+    }
+  }
+}