Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit 8158e4d

Browse files
kongweihankongweihan
committed
Add CSM for batch write flow control
1 parent f142db8 commit 8158e4d

7 files changed

Lines changed: 154 additions & 18 deletions

File tree

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

Lines changed: 43 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import com.google.bigtable.v2.MutateRowsResponse;
2727
import com.google.bigtable.v2.RateLimitInfo;
2828
import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
29+
import com.google.cloud.bigtable.data.v2.stub.metrics.Util;
2930
import com.google.common.annotations.VisibleForTesting;
3031
import com.google.common.base.Preconditions;
3132
import com.google.common.base.Stopwatch;
@@ -37,6 +38,9 @@
3738
import java.util.concurrent.atomic.AtomicReference;
3839
import java.util.logging.Logger;
3940
import javax.annotation.Nonnull;
41+
import javax.annotation.Nullable;
42+
43+
import static com.google.cloud.bigtable.data.v2.stub.metrics.Util.extractStatus;
4044

4145
class RateLimitingServerStreamingCallable
4246
extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
@@ -69,6 +73,8 @@ class RateLimitingServerStreamingCallable
6973

7074
private final ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable;
7175

76+
private BigtableTracer bigtableTracer;
77+
7278
RateLimitingServerStreamingCallable(
7379
@Nonnull ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable) {
7480
this.limiter = new ConditionalRateLimiter(DEFAULT_QPS);
@@ -84,8 +90,8 @@ public void call(
8490
limiter.acquire();
8591
stopwatch.stop();
8692
if (context.getTracer() instanceof BigtableTracer) {
87-
((BigtableTracer) context.getTracer())
88-
.batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
93+
bigtableTracer = (BigtableTracer) context.getTracer();
94+
bigtableTracer.batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
8995
}
9096
RateLimitingResponseObserver innerObserver = new RateLimitingResponseObserver(responseObserver);
9197
innerCallable.call(request, innerObserver, context);
@@ -104,7 +110,7 @@ static class ConditionalRateLimiter {
104110

105111
public ConditionalRateLimiter(long defaultQps) {
106112
limiter = RateLimiter.create(defaultQps);
107-
logger.info("Rate limiting is initiated (but disabled) with rate of " + defaultQps + " QPS.");
113+
logger.info("Batch write flow control: rate limiter is initiated (but disabled) with rate of " + defaultQps + " QPS.");
108114
}
109115

110116
/**
@@ -128,7 +134,7 @@ public void tryDisable() {
128134
if (now.isAfter(nextTime)) {
129135
boolean wasEnabled = this.enabled.getAndSet(false);
130136
if (wasEnabled) {
131-
logger.info("Rate limiter is disabled.");
137+
logger.info("Batch write flow control: rate limiter is disabled.");
132138
}
133139
// No need to update nextRateUpdateTime, any new RateLimitInfo can enable rate limiting and
134140
// update the rate again.
@@ -139,7 +145,7 @@ public void tryDisable() {
139145
public void enable() {
140146
boolean wasEnabled = this.enabled.getAndSet(true);
141147
if (!wasEnabled) {
142-
logger.info("Rate limiter is enabled.");
148+
logger.info("Batch write flow control: rate limiter is enabled.");
143149
}
144150
}
145151

@@ -158,31 +164,49 @@ public double getRate() {
158164
* @param rate The new rate of the rate limiter.
159165
* @param period The period during which rate should not be updated again and the rate limiter
160166
* should not be disabled.
167+
* @param bigtableTracer The tracer for exporting client-side metrics.
161168
*/
162-
public void trySetRate(double rate, Duration period) {
169+
public void trySetRate(
170+
double rate,
171+
Duration period,
172+
BigtableTracer bigtableTracer,
173+
double factor,
174+
String statusString) {
163175
Instant nextTime = nextRateUpdateTime.get();
164176
Instant now = Instant.now();
165177

166178
if (now.isBefore(nextTime)) {
179+
if (bigtableTracer != null) {
180+
bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, false);
181+
}
167182
return;
168183
}
169184

170185
Instant newNextTime = now.plusSeconds(period.getSeconds());
171186

172187
if (!nextRateUpdateTime.compareAndSet(nextTime, newNextTime)) {
173188
// Someone else updated it already.
189+
if (bigtableTracer != null) {
190+
bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, false);
191+
}
174192
return;
175193
}
176194
final double oldRate = limiter.getRate();
177195
limiter.setRate(rate);
178196
logger.info(
179-
"Updated max rate from "
197+
"Batch write flow control: updated max rate from "
180198
+ oldRate
181199
+ " to "
182200
+ rate
201+
+ " applied factor "
202+
+ factor
183203
+ " with period "
184204
+ period.getSeconds()
185205
+ " seconds.");
206+
if (bigtableTracer != null) {
207+
bigtableTracer.setBatchWriteFlowControlTargetQps(rate);
208+
bigtableTracer.addBatchWriteFlowControlFactor(factor, statusString, true);
209+
}
186210
}
187211

188212
@VisibleForTesting
@@ -215,17 +239,17 @@ private boolean hasValidRateLimitInfo(MutateRowsResponse response) {
215239
// have presence even thought it's marked as "optional". Check the factor and
216240
// period to make sure they're not 0.
217241
if (!response.hasRateLimitInfo()) {
218-
logger.finest("Response carries no RateLimitInfo");
242+
logger.finest("Batch write flow control: response carries no RateLimitInfo");
219243
return false;
220244
}
221245

222246
if (response.getRateLimitInfo().getFactor() <= 0
223247
|| response.getRateLimitInfo().getPeriod().getSeconds() <= 0) {
224-
logger.finest("Response carries invalid RateLimitInfo=" + response.getRateLimitInfo());
248+
logger.finest("Batch write flow control: response carries invalid RateLimitInfo=" + response.getRateLimitInfo());
225249
return false;
226250
}
227251

228-
logger.finest("Response carries valid RateLimitInfo=" + response.getRateLimitInfo());
252+
logger.finest("Batch write flow control: response carries valid RateLimitInfo=" + response.getRateLimitInfo());
229253
return true;
230254
}
231255

@@ -236,7 +260,8 @@ protected void onResponseImpl(MutateRowsResponse response) {
236260
RateLimitInfo info = response.getRateLimitInfo();
237261
updateQps(
238262
info.getFactor(),
239-
Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
263+
Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())),
264+
extractStatus(null));
240265
} else {
241266
limiter.tryDisable();
242267
}
@@ -250,7 +275,11 @@ protected void onErrorImpl(Throwable t) {
250275
if (t instanceof DeadlineExceededException
251276
|| t instanceof UnavailableException
252277
|| t instanceof ResourceExhaustedException) {
253-
updateQps(MIN_FACTOR, DEFAULT_PERIOD);
278+
logger.info("Batch write flow control: received error "
279+
+ extractStatus(t) + " applying min factor "
280+
+ MIN_FACTOR + " with period "
281+
+ DEFAULT_PERIOD + " seconds.");
282+
updateQps(MIN_FACTOR, DEFAULT_PERIOD, extractStatus(t));
254283
}
255284
outerObserver.onError(t);
256285
}
@@ -260,11 +289,11 @@ protected void onCompleteImpl() {
260289
outerObserver.onComplete();
261290
}
262291

263-
private void updateQps(double factor, Duration period) {
292+
private void updateQps(double factor, Duration period, String statusString) {
264293
double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR);
265294
double currentRate = limiter.getRate();
266295
double cappedRate = Math.min(Math.max(currentRate * cappedFactor, MIN_QPS), MAX_QPS);
267-
limiter.trySetRate(cappedRate, period);
296+
limiter.trySetRate(cappedRate, period, bigtableTracer, cappedFactor, statusString);
268297
}
269298
}
270299

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package com.google.cloud.bigtable.data.v2.stub.metrics;
1717

1818
import com.google.api.core.BetaApi;
19+
import com.google.api.core.InternalApi;
1920
import com.google.api.gax.rpc.ApiCallContext;
2021
import com.google.api.gax.tracing.ApiTracer;
2122
import com.google.api.gax.tracing.BaseApiTracer;
@@ -115,4 +116,25 @@ public void grpcMessageSent() {
115116
public void setTotalTimeoutDuration(Duration totalTimeoutDuration) {
116117
// noop
117118
}
119+
120+
/**
121+
* Record the target QPS for batch write flow control.
122+
*
123+
* @param targetQps The new target QPS for the client.
124+
*/
125+
@InternalApi
126+
public void setBatchWriteFlowControlTargetQps(double targetQps) {}
127+
128+
/**
129+
* Record the factors received from server-side for batch write flow control. The factors are
130+
* capped by min and max allowed factor values. Status and whether the factor was actually applied
131+
* are also recorded.
132+
*
133+
* @param factor Capped factor from server-side. For non-OK response, min factor is used.
134+
* @param statusString Status code of the request.
135+
* @param applied Whether the factor was actually applied.
136+
*/
137+
@InternalApi
138+
public void addBatchWriteFlowControlFactor(
139+
double factor, String statusString, boolean applied) {}
118140
}

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class BuiltinMetricsConstants {
4949
static final AttributeKey<String> METHOD_KEY = AttributeKey.stringKey("method");
5050
static final AttributeKey<String> STATUS_KEY = AttributeKey.stringKey("status");
5151
static final AttributeKey<String> CLIENT_UID_KEY = AttributeKey.stringKey("client_uid");
52+
static final AttributeKey<Boolean> APPLIED_KEY = AttributeKey.booleanKey("applied");
5253

5354
static final AttributeKey<String> TRANSPORT_TYPE = AttributeKey.stringKey("transport_type");
5455
static final AttributeKey<String> TRANSPORT_REGION = AttributeKey.stringKey("transport_region");
@@ -95,6 +96,9 @@ public class BuiltinMetricsConstants {
9596
static final String CLIENT_BLOCKING_LATENCIES_NAME = "throttling_latencies";
9697
static final String PER_CONNECTION_ERROR_COUNT_NAME = "per_connection_error_count";
9798
static final String OUTSTANDING_RPCS_PER_CHANNEL_NAME = "connection_pool/outstanding_rpcs";
99+
static final String BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME =
100+
"batch_write_flow_control_target_qps";
101+
static final String BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME = "batch_write_flow_control_factor";
98102

99103
// Start allow list of metrics that will be exported as internal
100104
public static final Map<String, Set<String>> GRPC_METRICS =
@@ -210,6 +214,8 @@ public class BuiltinMetricsConstants {
210214
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,
211215
135.0, 140.0, 145.0, 150.0, 155.0, 160.0, 165.0, 170.0, 175.0, 180.0, 185.0, 190.0,
212216
195.0, 200.0));
217+
private static final Aggregation AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM =
218+
Aggregation.explicitBucketHistogram(ImmutableList.of(0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3));
213219

214220
static final Set<AttributeKey> COMMON_ATTRIBUTES =
215221
ImmutableSet.of(
@@ -367,7 +373,23 @@ public static Map<InstrumentSelector, View> getAllViews() {
367373
.addAll(COMMON_ATTRIBUTES)
368374
.add(STREAMING_KEY, STATUS_KEY)
369375
.build());
370-
376+
defineView(
377+
views,
378+
BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME,
379+
Aggregation.sum(),
380+
InstrumentType.GAUGE,
381+
"1",
382+
ImmutableSet.<AttributeKey>builder().addAll(COMMON_ATTRIBUTES).build());
383+
defineView(
384+
views,
385+
BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME,
386+
AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM,
387+
InstrumentType.HISTOGRAM,
388+
"1",
389+
ImmutableSet.<AttributeKey>builder()
390+
.addAll(COMMON_ATTRIBUTES)
391+
.add(STATUS_KEY, APPLIED_KEY)
392+
.build());
371393
return views.build();
372394
}
373395
}

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

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import static com.google.api.gax.tracing.ApiTracerFactory.OperationType;
1919
import static com.google.api.gax.util.TimeConversionUtils.toJavaTimeDuration;
20+
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.APPLIED_KEY;
2021
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CLIENT_NAME_KEY;
2122
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CLUSTER_ID_KEY;
2223
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.METHOD_KEY;
@@ -41,6 +42,7 @@
4142
import com.google.gson.reflect.TypeToken;
4243
import io.grpc.Deadline;
4344
import io.opentelemetry.api.common.Attributes;
45+
import io.opentelemetry.api.metrics.DoubleGauge;
4446
import io.opentelemetry.api.metrics.DoubleHistogram;
4547
import io.opentelemetry.api.metrics.LongCounter;
4648
import java.time.Duration;
@@ -136,6 +138,8 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
136138
private final DoubleHistogram remainingDeadlineHistogram;
137139
private final LongCounter connectivityErrorCounter;
138140
private final LongCounter retryCounter;
141+
private final DoubleGauge batchWriteFlowControlTargetQps;
142+
private final DoubleHistogram batchWriteFlowControlFactorHistogram;
139143

140144
BuiltinMetricsTracer(
141145
OperationType operationType,
@@ -150,7 +154,9 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
150154
DoubleHistogram applicationBlockingLatenciesHistogram,
151155
DoubleHistogram deadlineHistogram,
152156
LongCounter connectivityErrorCounter,
153-
LongCounter retryCounter) {
157+
LongCounter retryCounter,
158+
DoubleGauge batchWriteFlowControlTargetQps,
159+
DoubleHistogram batchWriteFlowControlFactorHistogram) {
154160
this.operationType = operationType;
155161
this.spanName = spanName;
156162
this.baseAttributes = attributes;
@@ -165,6 +171,8 @@ static TransportAttrs create(@Nullable String locality, @Nullable String backend
165171
this.remainingDeadlineHistogram = deadlineHistogram;
166172
this.connectivityErrorCounter = connectivityErrorCounter;
167173
this.retryCounter = retryCounter;
174+
this.batchWriteFlowControlTargetQps = batchWriteFlowControlTargetQps;
175+
this.batchWriteFlowControlFactorHistogram = batchWriteFlowControlFactorHistogram;
168176
}
169177

170178
@Override
@@ -496,4 +504,24 @@ private static double convertToMs(long nanoSeconds) {
496504
double toMs = 1e-6;
497505
return nanoSeconds * toMs;
498506
}
507+
508+
@Override
509+
public void setBatchWriteFlowControlTargetQps(double targetQps) {
510+
Attributes attributes = baseAttributes.toBuilder().put(METHOD_KEY, spanName.toString()).build();
511+
512+
batchWriteFlowControlTargetQps.set(targetQps, attributes);
513+
}
514+
515+
@Override
516+
public void addBatchWriteFlowControlFactor(
517+
double factor, String statusString, boolean applied) {
518+
Attributes attributes =
519+
baseAttributes.toBuilder()
520+
.put(METHOD_KEY, spanName.toString())
521+
.put(STATUS_KEY, statusString)
522+
.put(APPLIED_KEY, applied)
523+
.build();
524+
525+
batchWriteFlowControlFactorHistogram.record(factor, attributes);
526+
}
499527
}

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.APPLICATION_BLOCKING_LATENCIES_NAME;
1919
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.ATTEMPT_LATENCIES2_NAME;
2020
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.ATTEMPT_LATENCIES_NAME;
21+
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME;
22+
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME;
2123
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CLIENT_BLOCKING_LATENCIES_NAME;
2224
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.CONNECTIVITY_ERROR_COUNT_NAME;
2325
import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.FIRST_RESPONSE_LATENCIES_NAME;
@@ -34,6 +36,7 @@
3436
import com.google.api.gax.tracing.SpanName;
3537
import io.opentelemetry.api.OpenTelemetry;
3638
import io.opentelemetry.api.common.Attributes;
39+
import io.opentelemetry.api.metrics.DoubleGauge;
3740
import io.opentelemetry.api.metrics.DoubleHistogram;
3841
import io.opentelemetry.api.metrics.LongCounter;
3942
import io.opentelemetry.api.metrics.Meter;
@@ -61,6 +64,8 @@ public class BuiltinMetricsTracerFactory extends BaseApiTracerFactory {
6164
private final DoubleHistogram remainingDeadlineHistogram;
6265
private final LongCounter connectivityErrorCounter;
6366
private final LongCounter retryCounter;
67+
private final DoubleGauge batchWriteFlowControlTargetQps;
68+
private final DoubleHistogram batchWriteFlowControlFactorHistogram;
6469

6570
public static BuiltinMetricsTracerFactory create(
6671
OpenTelemetry openTelemetry, Attributes attributes) throws IOException {
@@ -147,6 +152,19 @@ public static BuiltinMetricsTracerFactory create(
147152
.setDescription("The number of additional RPCs sent after the initial attempt.")
148153
.setUnit(COUNT)
149154
.build();
155+
batchWriteFlowControlTargetQps =
156+
meter
157+
.gaugeBuilder(BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME)
158+
.setDescription("The current target QPS of the client under batch write flow control.")
159+
.setUnit("1")
160+
.build();
161+
batchWriteFlowControlFactorHistogram =
162+
meter
163+
.histogramBuilder(BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME)
164+
.setDescription(
165+
"The distribution of batch write flow control factors received from the server.")
166+
.setUnit("1")
167+
.build();
150168
}
151169

152170
@Override
@@ -164,6 +182,8 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op
164182
applicationBlockingLatenciesHistogram,
165183
remainingDeadlineHistogram,
166184
connectivityErrorCounter,
167-
retryCounter);
185+
retryCounter,
186+
batchWriteFlowControlTargetQps,
187+
batchWriteFlowControlFactorHistogram);
168188
}
169189
}

0 commit comments

Comments
 (0)