Add CSM for batch write flow control

Author: kongweihan

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

@@ -37,6 +37,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Logger;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 class RateLimitingServerStreamingCallable
     extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
@@ -69,6 +70,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 +87,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);
@@ -155,22 +158,26 @@ public double getRate() {
      * Sets the rate and the next rate update time based on period, if the current time exceeds the
      * next rate update time. Otherwise, no-op.
      *
-     * @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 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 cappedFactor,
+                           @Nullable Throwable status) {
       Instant nextTime = nextRateUpdateTime.get();
       Instant now = Instant.now();
 
       if (now.isBefore(nextTime)) {
+        bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, false);
         return;
       }
 
       Instant newNextTime = now.plusSeconds(period.getSeconds());
 
       if (!nextRateUpdateTime.compareAndSet(nextTime, newNextTime)) {
         // Someone else updated it already.
+        bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, false);
         return;
       }
       final double oldRate = limiter.getRate();
@@ -183,6 +190,8 @@ public void trySetRate(double rate, Duration period) {
               + " with period "
               + period.getSeconds()
               + " seconds.");
+      bigtableTracer.setBatchWriteFlowControlTargetQps(rate);
+      bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, true);
     }
 
     @VisibleForTesting
@@ -236,7 +245,7 @@ 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())), null);
       } else {
         limiter.tryDisable();
       }
@@ -250,7 +259,7 @@ protected void onErrorImpl(Throwable t) {
       if (t instanceof DeadlineExceededException
           || t instanceof UnavailableException
           || t instanceof ResourceExhaustedException) {
-        updateQps(MIN_FACTOR, DEFAULT_PERIOD);
+        updateQps(MIN_FACTOR, DEFAULT_PERIOD, t);
       }
       outerObserver.onError(t);
     }
@@ -260,11 +269,11 @@ protected void onCompleteImpl() {
       outerObserver.onComplete();
     }
 
-    private void updateQps(double factor, Duration period) {
+    private void updateQps(double factor, Duration period, @Nullable Throwable t) {
       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, t);
     }
   }
 

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

@@ -41,9 +41,7 @@ public void onRequest(int requestCount) {
     // noop
   }
 
-  /**
-   * annotate when automatic flow control is disabled. This will be called in BuiltinMetricsTracer.
-   */
+  /** annotate when automatic flow control is disabled. This will be called in BuiltinMetricsTracer. */
   public void disableFlowControl() {
     // noop
   }
@@ -92,7 +90,8 @@ public void setLocations(String zone, String cluster) {
   }
 
   /** Set the underlying transport used to process the attempt */
-  public void setTransportAttrs(BuiltinMetricsTracer.TransportAttrs attrs) {}
+  public void setTransportAttrs(BuiltinMetricsTracer.TransportAttrs attrs) {
+  }
 
   @Deprecated
   /**
@@ -115,4 +114,23 @@ 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.
+   */
+  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 status  Status of the request.
+   * @param applied Whether the factor was actually applied.
+   */
+  public void addBatchWriteFlowControlFactor(double factor, @Nullable Throwable status, boolean applied) {
+  }
 }

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

@@ -49,6 +49,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");
@@ -70,6 +71,8 @@ public class BuiltinMetricsConstants {
   static final String REMAINING_DEADLINE_NAME = "remaining_deadline";
   static final String CLIENT_BLOCKING_LATENCIES_NAME = "throttling_latencies";
   static final String PER_CONNECTION_ERROR_COUNT_NAME = "per_connection_error_count";
+  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 =
@@ -140,6 +143,11 @@ public class BuiltinMetricsConstants {
               500_000.0,
               1_000_000.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(
           BIGTABLE_PROJECT_ID_KEY,
@@ -286,7 +294,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,
+        Aggregation.sum(),
+        InstrumentType.GAUGE,
+        "1",
+        ImmutableSet.<AttributeKey>builder().addAll(COMMON_ATTRIBUTES).add(STATUS_KEY).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(STREAMING_KEY, STATUS_KEY)
+            .build());
     return views.build();
   }
 }

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

@@ -28,6 +28,7 @@
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.TRANSPORT_TYPE;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.TRANSPORT_ZONE;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.ZONE_ID_KEY;
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.APPLIED_KEY;
 
 import com.google.api.core.ObsoleteApi;
 import com.google.api.gax.retrying.ServerStreamingAttemptException;
@@ -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,28 @@ 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, @Nullable Throwable status, boolean applied) {
+    String statusStr = status == null ? "OK" : Util.extractStatus(status);
+
+    Attributes attributes =
+        baseAttributes.toBuilder()
+            .put(METHOD_KEY, spanName.toString())
+            .put(STATUS_KEY, statusStr)
+            .put(APPLIED_KEY, applied)
+            .build();
+
+    batchWriteFlowControlFactorHistogram.record(factor, attributes);
+  }
 }

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

@@ -26,6 +26,8 @@
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.REMAINING_DEADLINE_NAME;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.RETRY_COUNT_NAME;
 import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.SERVER_LATENCIES_NAME;
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME;
+import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME;
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.tracing.ApiTracer;
@@ -34,9 +36,11 @@
 import com.google.api.gax.tracing.SpanName;
 import io.opentelemetry.api.OpenTelemetry;
 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 io.opentelemetry.api.metrics.Meter;
+
 import java.io.IOException;
 
 /**
@@ -61,6 +65,8 @@ public class BuiltinMetricsTracerFactory extends BaseApiTracerFactory {
   private final DoubleHistogram remainingDeadlineHistogram;
   private final LongCounter connectivityErrorCounter;
   private final LongCounter retryCounter;
+  private final DoubleGauge batchWriteFlowControlTargetQps;
+  private final DoubleHistogram batchWriteFlowControlFactorHistogram;
 
   public static BuiltinMetricsTracerFactory create(
       OpenTelemetry openTelemetry, Attributes attributes) throws IOException {
@@ -147,6 +153,18 @@ public static BuiltinMetricsTracerFactory create(
             .setDescription("The number of additional RPCs sent after the initial attempt.")
             .setUnit(COUNT)
             .build();
+    batchWriteFlowControlTargetQps =
+        meter
+            .gaugeBuilder(BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME)
+            .setDescription("The current target QPS of the client under batch write flow control.")
+            .setUnit("1").build();
+    batchWriteFlowControlFactorHistogram =
+        meter
+            .histogramBuilder(BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME)
+            .setDescription(
+                "The distribution of batch write flow control factors received from the server.")
+            .setUnit("1")
+            .build();
   }
 
   @Override
@@ -164,6 +182,8 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op
         applicationBlockingLatenciesHistogram,
         remainingDeadlineHistogram,
         connectivityErrorCounter,
-        retryCounter);
+        retryCounter,
+        batchWriteFlowControlTargetQps,
+        batchWriteFlowControlFactorHistogram);
   }
 }

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

@@ -266,4 +266,18 @@ public void setTotalTimeoutDuration(java.time.Duration totalTimeoutDuration) {
       tracer.setTotalTimeoutDuration(totalTimeoutDuration);
     }
   }
+
+  @Override
+  public void setBatchWriteFlowControlTargetQps(double targetQps) {
+    for (BigtableTracer tracer : bigtableTracers) {
+      tracer.setBatchWriteFlowControlTargetQps(targetQps);
+    }
+  }
+
+  @Override
+  public void addBatchWriteFlowControlFactor(double factor, @Nullable Throwable t, boolean applied) {
+    for (BigtableTracer tracer : bigtableTracers) {
+      tracer.addBatchWriteFlowControlFactor(factor, t, applied);
+    }
+  }
 }