Skip to content

Commit e8a246f

Browse files
kongweihankongweihan
committed
Add CSM for batch write flow control
1 parent 6998e3b commit e8a246f

File tree

6 files changed

+131
-10
lines changed

6 files changed

+131
-10
lines changed

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

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.util.concurrent.atomic.AtomicReference;
3838
import java.util.logging.Logger;
3939
import javax.annotation.Nonnull;
40+
import javax.annotation.Nullable;
4041

4142
class RateLimitingServerStreamingCallable
4243
extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
@@ -69,6 +70,8 @@ class RateLimitingServerStreamingCallable
6970

7071
private final ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable;
7172

73+
private BigtableTracer bigtableTracer;
74+
7275
RateLimitingServerStreamingCallable(
7376
@Nonnull ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable) {
7477
this.limiter = new ConditionalRateLimiter(DEFAULT_QPS);
@@ -84,8 +87,8 @@ public void call(
8487
limiter.acquire();
8588
stopwatch.stop();
8689
if (context.getTracer() instanceof BigtableTracer) {
87-
((BigtableTracer) context.getTracer())
88-
.batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
90+
bigtableTracer = (BigtableTracer) context.getTracer();
91+
bigtableTracer.batchRequestThrottled(stopwatch.elapsed(TimeUnit.NANOSECONDS));
8992
}
9093
RateLimitingResponseObserver innerObserver = new RateLimitingResponseObserver(responseObserver);
9194
innerCallable.call(request, innerObserver, context);
@@ -158,19 +161,27 @@ public double getRate() {
158161
* @param rate The new rate of the rate limiter.
159162
* @param period The period during which rate should not be updated again and the rate limiter
160163
* should not be disabled.
164+
* @param bigtableTracer The tracer for exporting client-side metrics.
161165
*/
162-
public void trySetRate(double rate, Duration period) {
166+
public void trySetRate(
167+
double rate,
168+
Duration period,
169+
BigtableTracer bigtableTracer,
170+
double cappedFactor,
171+
@Nullable Throwable status) {
163172
Instant nextTime = nextRateUpdateTime.get();
164173
Instant now = Instant.now();
165174

166175
if (now.isBefore(nextTime)) {
176+
bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, false);
167177
return;
168178
}
169179

170180
Instant newNextTime = now.plusSeconds(period.getSeconds());
171181

172182
if (!nextRateUpdateTime.compareAndSet(nextTime, newNextTime)) {
173183
// Someone else updated it already.
184+
bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, false);
174185
return;
175186
}
176187
final double oldRate = limiter.getRate();
@@ -183,6 +194,8 @@ public void trySetRate(double rate, Duration period) {
183194
+ " with period "
184195
+ period.getSeconds()
185196
+ " seconds.");
197+
bigtableTracer.setBatchWriteFlowControlTargetQps(rate);
198+
bigtableTracer.addBatchWriteFlowControlFactor(cappedFactor, status, true);
186199
}
187200

188201
@VisibleForTesting
@@ -236,7 +249,8 @@ protected void onResponseImpl(MutateRowsResponse response) {
236249
RateLimitInfo info = response.getRateLimitInfo();
237250
updateQps(
238251
info.getFactor(),
239-
Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
252+
Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())),
253+
null);
240254
} else {
241255
limiter.tryDisable();
242256
}
@@ -250,7 +264,7 @@ protected void onErrorImpl(Throwable t) {
250264
if (t instanceof DeadlineExceededException
251265
|| t instanceof UnavailableException
252266
|| t instanceof ResourceExhaustedException) {
253-
updateQps(MIN_FACTOR, DEFAULT_PERIOD);
267+
updateQps(MIN_FACTOR, DEFAULT_PERIOD, t);
254268
}
255269
outerObserver.onError(t);
256270
}
@@ -260,11 +274,11 @@ protected void onCompleteImpl() {
260274
outerObserver.onComplete();
261275
}
262276

263-
private void updateQps(double factor, Duration period) {
277+
private void updateQps(double factor, Duration period, @Nullable Throwable t) {
264278
double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR);
265279
double currentRate = limiter.getRate();
266280
double cappedRate = Math.min(Math.max(currentRate * cappedFactor, MIN_QPS), MAX_QPS);
267-
limiter.trySetRate(cappedRate, period);
281+
limiter.trySetRate(cappedRate, period, bigtableTracer, cappedFactor, t);
268282
}
269283
}
270284

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,4 +115,23 @@ public void grpcMessageSent() {
115115
public void setTotalTimeoutDuration(Duration totalTimeoutDuration) {
116116
// noop
117117
}
118+
119+
/**
120+
* Record the target QPS for batch write flow control.
121+
*
122+
* @param targetQps The new target QPS for the client.
123+
*/
124+
public void setBatchWriteFlowControlTargetQps(double targetQps) {}
125+
126+
/**
127+
* Record the factors received from server-side for batch write flow control. The factors are
128+
* capped by min and max allowed factor values. Status and whether the factor was actually applied
129+
* are also recorded.
130+
*
131+
* @param factor Capped factor from server-side. For non-OK response, min factor is used.
132+
* @param status Status of the request.
133+
* @param applied Whether the factor was actually applied.
134+
*/
135+
public void addBatchWriteFlowControlFactor(
136+
double factor, @Nullable Throwable status, boolean applied) {}
118137
}

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

Lines changed: 24 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");
@@ -70,6 +71,9 @@ public class BuiltinMetricsConstants {
7071
static final String REMAINING_DEADLINE_NAME = "remaining_deadline";
7172
static final String CLIENT_BLOCKING_LATENCIES_NAME = "throttling_latencies";
7273
static final String PER_CONNECTION_ERROR_COUNT_NAME = "per_connection_error_count";
74+
static final String BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME =
75+
"batch_write_flow_control_target_qps";
76+
static final String BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME = "batch_write_flow_control_factor";
7377

7478
// Start allow list of metrics that will be exported as internal
7579
public static final Map<String, Set<String>> GRPC_METRICS =
@@ -140,6 +144,9 @@ public class BuiltinMetricsConstants {
140144
500_000.0,
141145
1_000_000.0));
142146

147+
private static final Aggregation AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM =
148+
Aggregation.explicitBucketHistogram(ImmutableList.of(0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3));
149+
143150
static final Set<AttributeKey> COMMON_ATTRIBUTES =
144151
ImmutableSet.of(
145152
BIGTABLE_PROJECT_ID_KEY,
@@ -286,7 +293,23 @@ public static Map<InstrumentSelector, View> getAllViews() {
286293
.addAll(COMMON_ATTRIBUTES)
287294
.add(STREAMING_KEY, STATUS_KEY)
288295
.build());
289-
296+
defineView(
297+
views,
298+
BATCH_WRITE_FLOW_CONTROL_TARGET_QPS_NAME,
299+
Aggregation.sum(),
300+
InstrumentType.GAUGE,
301+
"1",
302+
ImmutableSet.<AttributeKey>builder().addAll(COMMON_ATTRIBUTES).add(STATUS_KEY).build());
303+
defineView(
304+
views,
305+
BATCH_WRITE_FLOW_CONTROL_FACTOR_NAME,
306+
AGGREGATION_BATCH_WRITE_FLOW_CONTROL_FACTOR_HISTOGRAM,
307+
InstrumentType.HISTOGRAM,
308+
"1",
309+
ImmutableSet.<AttributeKey>builder()
310+
.addAll(COMMON_ATTRIBUTES)
311+
.add(STREAMING_KEY, STATUS_KEY)
312+
.build());
290313
return views.build();
291314
}
292315
}

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

Lines changed: 31 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,26 @@ 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, @Nullable Throwable status, boolean applied) {
518+
String statusStr = status == null ? "OK" : Util.extractStatus(status);
519+
520+
Attributes attributes =
521+
baseAttributes.toBuilder()
522+
.put(METHOD_KEY, spanName.toString())
523+
.put(STATUS_KEY, statusStr)
524+
.put(APPLIED_KEY, applied)
525+
.build();
526+
527+
batchWriteFlowControlFactorHistogram.record(factor, attributes);
528+
}
499529
}

‎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
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,4 +266,19 @@ public void setTotalTimeoutDuration(java.time.Duration totalTimeoutDuration) {
266266
tracer.setTotalTimeoutDuration(totalTimeoutDuration);
267267
}
268268
}
269+
270+
@Override
271+
public void setBatchWriteFlowControlTargetQps(double targetQps) {
272+
for (BigtableTracer tracer : bigtableTracers) {
273+
tracer.setBatchWriteFlowControlTargetQps(targetQps);
274+
}
275+
}
276+
277+
@Override
278+
public void addBatchWriteFlowControlFactor(
279+
double factor, @Nullable Throwable t, boolean applied) {
280+
for (BigtableTracer tracer : bigtableTracers) {
281+
tracer.addBatchWriteFlowControlFactor(factor, t, applied);
282+
}
283+
}
269284
}

0 commit comments

Comments
 (0)