Add CSM for batch write flow control

Author: kongweihan

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

@@ -15,6 +15,8 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
+import static com.google.cloud.bigtable.data.v2.stub.metrics.Util.extractStatus;
+
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.DeadlineExceededException;
 import com.google.api.gax.rpc.ResourceExhaustedException;
@@ -69,6 +71,8 @@ class RateLimitingServerStreamingCallable
 
   private final ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable;
 
+  private BigtableTracer bigtableTracer;
+
   RateLimitingServerStreamingCallable(
       @Nonnull ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable) {
     this.limiter = new ConditionalRateLimiter(DEFAULT_QPS);
@@ -84,8 +88,8 @@ public void call(
     limiter.acquire();
     stopwatch.stop();
     if (context.getTracer() instanceof BigtableTracer) {
-      ((BigtableTracer) context.getTracer())
-          .batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
+      bigtableTracer = (BigtableTracer) context.getTracer();
+      bigtableTracer.batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
     }
     RateLimitingResponseObserver innerObserver = new RateLimitingResponseObserver(responseObserver);
     innerCallable.call(request, innerObserver, context);
@@ -104,7 +108,10 @@ static class ConditionalRateLimiter {
 
     public ConditionalRateLimiter(long defaultQps) {
       limiter = RateLimiter.create(defaultQps);
-      logger.info("Rate limiting is initiated (but disabled) with rate of " + defaultQps + " QPS.");
+      logger.info(
+          "Batch write flow control: rate limiter is initiated (but disabled) with rate of "
+              + defaultQps
+              + " QPS.");
     }
 
     /**
@@ -128,7 +135,7 @@ public void tryDisable() {
       if (now.isAfter(nextTime)) {
         boolean wasEnabled = this.enabled.getAndSet(false);
         if (wasEnabled) {
-          logger.info("Rate limiter is disabled.");
+          logger.info("Batch write flow control: rate limiter is disabled.");
         }
         // No need to update nextRateUpdateTime, any new RateLimitInfo can enable rate limiting and
         // update the rate again.
@@ -139,7 +146,7 @@ public void tryDisable() {
     public void enable() {
       boolean wasEnabled = this.enabled.getAndSet(true);
       if (!wasEnabled) {
-        logger.info("Rate limiter is enabled.");
+        logger.info("Batch write flow control: rate limiter is enabled.");
       }
     }
 
@@ -158,31 +165,49 @@ public double getRate() {
      * @param rate The new rate of the rate limiter.
      * @param period The period during which rate should not be updated again and the rate limiter
      *     should not be disabled.
+     * @param bigtableTracer The tracer for exporting client-side metrics.
      */
-    public void trySetRate(double rate, Duration period) {
+    public void trySetRate(
+        double rate,
+        Duration period,
+        BigtableTracer bigtableTracer,
+        double factor,
+        String statusString) {
       Instant nextTime = nextRateUpdateTime.get();
       Instant now = Instant.now();
 
       if (now.isBefore(nextTime)) {
+        if (bigtableTracer != null) {
+          bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, false);
+        }
         return;
       }
 
       Instant newNextTime = now.plusSeconds(period.getSeconds());
 
       if (!nextRateUpdateTime.compareAndSet(nextTime, newNextTime)) {
         // Someone else updated it already.
+        if (bigtableTracer != null) {
+          bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, false);
+        }
         return;
       }
       final double oldRate = limiter.getRate();
       limiter.setRate(rate);
       logger.info(
-          "Updated max rate from "
+          "Batch write flow control: updated max rate from "
               + oldRate
               + " to "
               + rate
+              + " applied factor "
+              + factor
               + " with period "
               + period.getSeconds()
               + " seconds.");
+      if (bigtableTracer != null) {
+        bigtableTracer.setBatchWriteFlowControlTargetQps(rate);
+        bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, true);
+      }
     }
 
     @VisibleForTesting
@@ -215,17 +240,21 @@ private boolean hasValidRateLimitInfo(MutateRowsResponse response) {
       // 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.finest("Response carries no RateLimitInfo");
+        logger.finest("Batch write flow control: response carries no RateLimitInfo");
         return false;
       }
 
       if (response.getRateLimitInfo().getFactor() <= 0
           || response.getRateLimitInfo().getPeriod().getSeconds() <= 0) {
-        logger.finest("Response carries invalid RateLimitInfo=" + response.getRateLimitInfo());
+        logger.finest(
+            "Batch write flow control: response carries invalid RateLimitInfo="
+                + response.getRateLimitInfo());
         return false;
       }
 
-      logger.finest("Response carries valid RateLimitInfo=" + response.getRateLimitInfo());
+      logger.finest(
+          "Batch write flow control: response carries valid RateLimitInfo="
+              + response.getRateLimitInfo());
       return true;
     }
 
@@ -236,7 +265,8 @@ protected void onResponseImpl(MutateRowsResponse response) {
         RateLimitInfo info = response.getRateLimitInfo();
         updateQps(
             info.getFactor(),
-            Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
+            Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())),
+            extractStatus(null));
       } else {
         limiter.tryDisable();
       }
@@ -250,7 +280,15 @@ protected void onErrorImpl(Throwable t) {
       if (t instanceof DeadlineExceededException
           || t instanceof UnavailableException
           || t instanceof ResourceExhaustedException) {
-        updateQps(MIN_FACTOR, DEFAULT_PERIOD);
+        logger.info(
+            "Batch write flow control: received error "
+                + extractStatus(t)
+                + " applying min factor "
+                + MIN_FACTOR
+                + " with period "
+                + DEFAULT_PERIOD
+                + " seconds.");
+        updateQps(MIN_FACTOR, DEFAULT_PERIOD, extractStatus(t));
       }
       outerObserver.onError(t);
     }
@@ -260,11 +298,11 @@ protected void onCompleteImpl() {
       outerObserver.onComplete();
     }
 
-    private void updateQps(double factor, Duration period) {
+    private void updateQps(double factor, Duration period, String statusString) {
       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);
+      limiter.trySetRate(cappedRate, period, bigtableTracer, cappedFactor, statusString);
     }
   }
 

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

@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.data.v2.stub.metrics;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.tracing.ApiTracer;
 import com.google.api.gax.tracing.BaseApiTracer;
@@ -115,4 +116,24 @@ public void grpcMessageSent() {
   public void setTotalTimeoutDuration(Duration totalTimeoutDuration) {
     // noop
   }
+
+  /**
+   * Record the target QPS for batch write flow control.
+   *
+   * @param targetQps The new target QPS for the client.
+   */
+  @InternalApi
+  public void setBatchWriteFlowControlTargetQps(double targetQps) {}
+
+  /**
+   * Record the factors received from server-side for batch write flow control. The factors are
+   * capped by min and max allowed factor values. Status and whether the factor was actually applied
+   * are also recorded.
+   *
+   * @param factor Capped factor from server-side. For non-OK response, min factor is used.
+   * @param statusString Status code of the request.
+   * @param applied Whether the factor was actually applied.
+   */
+  @InternalApi
+  public void addBatchWriteFlowControlFactor(double factor, String statusString, boolean applied) {}
 }

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

@@ -20,10 +20,9 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import io.opentelemetry.api.common.AttributeKey;
-import io.opentelemetry.sdk.metrics.Aggregation;
-import io.opentelemetry.sdk.metrics.InstrumentSelector;
-import io.opentelemetry.sdk.metrics.InstrumentType;
-import io.opentelemetry.sdk.metrics.View;
+import io.opentelemetry.sdk.metrics.*;
+
+import javax.annotation.Nullable;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -49,6 +48,7 @@ public class BuiltinMetricsConstants {
   static final AttributeKey<String> METHOD_KEY = AttributeKey.stringKey("method");
   static final AttributeKey<String> STATUS_KEY = AttributeKey.stringKey("status");
   static final AttributeKey<String> CLIENT_UID_KEY = AttributeKey.stringKey("client_uid");
+  static final AttributeKey<Boolean> APPLIED_KEY = AttributeKey.booleanKey("applied");
 
   static final AttributeKey<String> TRANSPORT_TYPE = AttributeKey.stringKey("transport_type");
   static final AttributeKey<String> TRANSPORT_REGION = AttributeKey.stringKey("transport_region");
@@ -95,6 +95,9 @@ public class BuiltinMetricsConstants {
   static final String CLIENT_BLOCKING_LATENCIES_NAME = "throttling_latencies";
   static final String PER_CONNECTION_ERROR_COUNT_NAME = "per_connection_error_count";
   static final String OUTSTANDING_RPCS_PER_CHANNEL_NAME = "connection_pool/outstanding_rpcs";
+  static final String BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME =
+      "batch_write_flow_control_target_qps";
+  static final String BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME = "batch_write_flow_control_factor";
 
   // Start allow list of metrics that will be exported as internal
   public static final Map<String, Set<String>> GRPC_METRICS =
@@ -210,6 +213,8 @@ public class BuiltinMetricsConstants {
               70.0, 75.0, 80.0, 85.0, 90.0, 95.0, 100.0, 105.0, 110.0, 115.0, 120.0, 125.0, 130.0,
               135.0, 140.0, 145.0, 150.0, 155.0, 160.0, 165.0, 170.0, 175.0, 180.0, 185.0, 190.0,
               195.0, 200.0));
+  private static final Aggregation AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM =
+      Aggregation.explicitBucketHistogram(ImmutableList.of(0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3));
 
   static final Set<AttributeKey> COMMON_ATTRIBUTES =
       ImmutableSet.of(
@@ -225,7 +230,7 @@ public class BuiltinMetricsConstants {
   static void defineView(
       ImmutableMap.Builder<InstrumentSelector, View> viewMap,
       String id,
-      Aggregation aggregation,
+      @Nullable Aggregation aggregation,
       InstrumentType type,
       String unit,
       Set<AttributeKey> attributes) {
@@ -242,14 +247,14 @@ static void defineView(
                 COMMON_ATTRIBUTES.stream().map(AttributeKey::getKey).collect(Collectors.toSet()))
             .addAll(attributes.stream().map(AttributeKey::getKey).collect(Collectors.toSet()))
             .build();
-    View view =
+    ViewBuilder viewBuilder =
         View.builder()
             .setName(METER_NAME + id)
-            .setAggregation(aggregation)
-            .setAttributeFilter(attributesFilter)
-            .build();
-
-    viewMap.put(selector, view);
+            .setAttributeFilter(attributesFilter);
+    if (aggregation != null) {
+      viewBuilder.setAggregation(aggregation);
+    }
+    viewMap.put(selector, viewBuilder.build());
   }
 
   // uses cloud.BigtableClient schema
@@ -367,7 +372,23 @@ public static Map<InstrumentSelector, View> getAllViews() {
             .addAll(COMMON_ATTRIBUTES)
             .add(STREAMING_KEY, STATUS_KEY)
             .build());
-
+    defineView(
+        views,
+        BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME,
+        null,
+        InstrumentType.GAUGE,
+        "1",
+        ImmutableSet.<AttributeKey>builder().addAll(COMMON_ATTRIBUTES).build());
+    defineView(
+        views,
+        BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME,
+        AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM,
+        InstrumentType.HISTOGRAM,
+        "1",
+        ImmutableSet.<AttributeKey>builder()
+            .addAll(COMMON_ATTRIBUTES)
+            .add(STATUS_KEY, APPLIED_KEY)
+            .build());
     return views.build();
   }
 }

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

@@ -17,6 +17,7 @@
 
 import static com.google.api.gax.tracing.ApiTracerFactory.OperationType;
 import static com.google.api.gax.util.TimeConversionUtils.toJavaTimeDuration;
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.APPLIED_KEY;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CLIENT_NAME_KEY;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CLUSTER_ID_KEY;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.METHOD_KEY;
@@ -41,6 +42,7 @@
 import com.google.gson.reflect.TypeToken;
 import io.grpc.Deadline;
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.metrics.DoubleGauge;
 import io.opentelemetry.api.metrics.DoubleHistogram;
 import io.opentelemetry.api.metrics.LongCounter;
 import java.time.Duration;
@@ -136,6 +138,8 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
   private final DoubleHistogram remainingDeadlineHistogram;
   private final LongCounter connectivityErrorCounter;
   private final LongCounter retryCounter;
+  private final DoubleGauge batchWriteFlowControlTargetQps;
+  private final DoubleHistogram batchWriteFlowControlFactorHistogram;
 
   BuiltinMetricsTracer(
       OperationType operationType,
@@ -150,7 +154,9 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
       DoubleHistogram applicationBlockingLatenciesHistogram,
       DoubleHistogram deadlineHistogram,
       LongCounter connectivityErrorCounter,
-      LongCounter retryCounter) {
+      LongCounter retryCounter,
+      DoubleGauge batchWriteFlowControlTargetQps,
+      DoubleHistogram batchWriteFlowControlFactorHistogram) {
     this.operationType = operationType;
     this.spanName = spanName;
     this.baseAttributes = attributes;
@@ -165,6 +171,8 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
     this.remainingDeadlineHistogram = deadlineHistogram;
     this.connectivityErrorCounter = connectivityErrorCounter;
     this.retryCounter = retryCounter;
+    this.batchWriteFlowControlTargetQps = batchWriteFlowControlTargetQps;
+    this.batchWriteFlowControlFactorHistogram = batchWriteFlowControlFactorHistogram;
   }
 
   @Override
@@ -496,4 +504,23 @@ private static double convertToMs(long nanoSeconds) {
     double toMs = 1e-6;
     return nanoSeconds * toMs;
   }
+
+  @Override
+  public void setBatchWriteFlowControlTargetQps(double targetQps) {
+    Attributes attributes = baseAttributes.toBuilder().put(METHOD_KEY, spanName.toString()).build();
+
+    batchWriteFlowControlTargetQps.set(targetQps, attributes);
+  }
+
+  @Override
+  public void addBatchWriteFlowControlFactor(double factor, String statusString, boolean applied) {
+    Attributes attributes =
+        baseAttributes.toBuilder()
+            .put(METHOD_KEY, spanName.toString())
+            .put(STATUS_KEY, statusString)
+            .put(APPLIED_KEY, applied)
+            .build();
+
+    batchWriteFlowControlFactorHistogram.record(factor, attributes);
+  }
 }