Skip to content

Commit ee0da11

Browse files
mutianfmutianf
authored
feat: add rate limiting callable (#1729)
* feat: add rate limiting callable * address comments
1 parent 7e5c646 commit ee0da11

File tree

2 files changed

+361
-0
lines changed

2 files changed

+361
-0
lines changed

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

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub;
17+
18+
import com.google.api.gax.rpc.ApiCallContext;
19+
import com.google.api.gax.rpc.DeadlineExceededException;
20+
import com.google.api.gax.rpc.ResourceExhaustedException;
21+
import com.google.api.gax.rpc.ResponseObserver;
22+
import com.google.api.gax.rpc.ServerStreamingCallable;
23+
import com.google.api.gax.rpc.StreamController;
24+
import com.google.api.gax.rpc.UnavailableException;
25+
import com.google.bigtable.v2.MutateRowsRequest;
26+
import com.google.bigtable.v2.MutateRowsResponse;
27+
import com.google.bigtable.v2.RateLimitInfo;
28+
import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
29+
import com.google.common.annotations.VisibleForTesting;
30+
import com.google.common.base.Preconditions;
31+
import com.google.common.base.Stopwatch;
32+
import com.google.common.util.concurrent.RateLimiter;
33+
import java.util.concurrent.TimeUnit;
34+
import java.util.concurrent.atomic.AtomicReference;
35+
import java.util.logging.Level;
36+
import java.util.logging.Logger;
37+
import javax.annotation.Nonnull;
38+
import org.threeten.bp.Duration;
39+
import org.threeten.bp.Instant;
40+
41+
class RateLimitingServerStreamingCallable
42+
extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
43+
private static final Logger logger =
44+
Logger.getLogger(RateLimitingServerStreamingCallable.class.getName());
45+
46+
// When the mutation size is large, starting with a higher QPS will make
47+
// the dataflow job fail very quickly. Start with lower QPS and increase
48+
// the QPS gradually if the server doesn't push back
49+
private static final long DEFAULT_QPS = 10;
50+
51+
// Default interval before changing the QPS on error responses
52+
private static final Duration DEFAULT_PERIOD = Duration.ofSeconds(10);
53+
54+
// Minimum QPS to make sure the job is not stuck
55+
private static final double MIN_QPS = 0.1;
56+
private static final double MAX_QPS = 100_000;
57+
58+
// QPS can be lowered to at most MIN_FACTOR * currentQps. When server returned
59+
// an error, use MIN_FACTOR to calculate the new QPS. This is the same as
60+
// the server side cap.
61+
@VisibleForTesting static final double MIN_FACTOR = 0.7;
62+
63+
// QPS can be increased to at most MAX_FACTOR * currentQps. This is the same
64+
// as the server side cap
65+
private static final double MAX_FACTOR = 1.3;
66+
67+
private final RateLimiter limiter;
68+
69+
private final AtomicReference<Instant> lastQpsChangeTime = new AtomicReference<>(Instant.now());
70+
private final ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable;
71+
72+
RateLimitingServerStreamingCallable(
73+
@Nonnull ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> innerCallable) {
74+
this.limiter = RateLimiter.create(DEFAULT_QPS);
75+
this.innerCallable = Preconditions.checkNotNull(innerCallable, "Inner callable must be set");
76+
logger.info("Rate limiting is enabled with initial QPS of " + limiter.getRate());
77+
}
78+
79+
@Override
80+
public void call(
81+
MutateRowsRequest request,
82+
ResponseObserver<MutateRowsResponse> responseObserver,
83+
ApiCallContext context) {
84+
Stopwatch stopwatch = Stopwatch.createStarted();
85+
limiter.acquire();
86+
stopwatch.stop();
87+
if (context.getTracer() instanceof BigtableTracer) {
88+
((BigtableTracer) context.getTracer())
89+
.batchRequestThrottled(stopwatch.elapsed(TimeUnit.MILLISECONDS));
90+
}
91+
RateLimitingResponseObserver innerObserver =
92+
new RateLimitingResponseObserver(limiter, lastQpsChangeTime, responseObserver);
93+
innerCallable.call(request, innerObserver, context);
94+
}
95+
96+
class RateLimitingResponseObserver extends SafeResponseObserver<MutateRowsResponse> {
97+
private final ResponseObserver<MutateRowsResponse> outerObserver;
98+
private final RateLimiter rateLimiter;
99+
100+
private final AtomicReference<Instant> lastQpsChangeTime;
101+
102+
RateLimitingResponseObserver(
103+
RateLimiter rateLimiter,
104+
AtomicReference<Instant> lastQpsChangeTime,
105+
ResponseObserver<MutateRowsResponse> observer) {
106+
super(observer);
107+
this.outerObserver = observer;
108+
this.rateLimiter = rateLimiter;
109+
this.lastQpsChangeTime = lastQpsChangeTime;
110+
}
111+
112+
@Override
113+
protected void onStartImpl(StreamController controller) {
114+
outerObserver.onStart(controller);
115+
}
116+
117+
@Override
118+
protected void onResponseImpl(MutateRowsResponse response) {
119+
if (response.hasRateLimitInfo()) {
120+
RateLimitInfo info = response.getRateLimitInfo();
121+
// RateLimitInfo is an optional field. However, proto3 sub-message field always
122+
// have presence even thought it's marked as "optional". Check the factor and
123+
// period to make sure they're not 0.
124+
if (info.getFactor() != 0 && info.getPeriod().getSeconds() != 0) {
125+
updateQps(
126+
info.getFactor(),
127+
Duration.ofSeconds(com.google.protobuf.util.Durations.toSeconds(info.getPeriod())));
128+
}
129+
}
130+
}
131+
132+
@Override
133+
protected void onErrorImpl(Throwable t) {
134+
// When server returns DEADLINE_EXCEEDED, UNAVAILABLE or RESOURCE_EXHAUSTED,
135+
// assume cbt server is overloaded
136+
if (t instanceof DeadlineExceededException
137+
|| t instanceof UnavailableException
138+
|| t instanceof ResourceExhaustedException) {
139+
updateQps(MIN_FACTOR, DEFAULT_PERIOD);
140+
}
141+
outerObserver.onError(t);
142+
}
143+
144+
@Override
145+
protected void onCompleteImpl() {
146+
outerObserver.onComplete();
147+
}
148+
149+
private void updateQps(double factor, Duration period) {
150+
Instant lastTime = lastQpsChangeTime.get();
151+
Instant now = Instant.now();
152+
153+
if (now.minus(period).isAfter(lastTime) && lastQpsChangeTime.compareAndSet(lastTime, now)) {
154+
double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR);
155+
double currentRate = limiter.getRate();
156+
limiter.setRate(Math.min(Math.max(currentRate * cappedFactor, MIN_QPS), MAX_QPS));
157+
logger.log(
158+
Level.FINE,
159+
"Updated QPS from {0} to {1}, server returned factor is {2}, capped factor is {3}",
160+
new Object[] {currentRate, limiter.getRate(), factor, cappedFactor});
161+
}
162+
}
163+
}
164+
165+
@VisibleForTesting
166+
AtomicReference<Instant> getLastQpsChangeTime() {
167+
return lastQpsChangeTime;
168+
}
169+
170+
@VisibleForTesting
171+
double getCurrentRate() {
172+
return limiter.getRate();
173+
}
174+
}

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

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigtable.data.v2.stub;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
21+
import com.google.api.gax.grpc.GrpcCallContext;
22+
import com.google.api.gax.rpc.ApiCallContext;
23+
import com.google.api.gax.rpc.DeadlineExceededException;
24+
import com.google.api.gax.rpc.ResponseObserver;
25+
import com.google.api.gax.rpc.ServerStreamingCallable;
26+
import com.google.api.gax.rpc.StatusCode;
27+
import com.google.api.gax.rpc.StreamController;
28+
import com.google.bigtable.v2.MutateRowsRequest;
29+
import com.google.bigtable.v2.MutateRowsResponse;
30+
import com.google.bigtable.v2.RateLimitInfo;
31+
import com.google.cloud.bigtable.gaxx.testing.FakeStatusCode;
32+
import com.google.protobuf.Duration;
33+
import org.junit.Before;
34+
import org.junit.Test;
35+
import org.junit.runner.RunWith;
36+
import org.junit.runners.JUnit4;
37+
import org.mockito.Mockito;
38+
import org.threeten.bp.Instant;
39+
40+
@RunWith(JUnit4.class)
41+
public class RateLimitingCallableTest {
42+
43+
private final MutateRowsRequest request =
44+
MutateRowsRequest.newBuilder().getDefaultInstanceForType();
45+
private final ResponseObserver<MutateRowsResponse> responseObserver =
46+
Mockito.mock(ResponseObserver.class);
47+
private final ApiCallContext context = GrpcCallContext.createDefault();
48+
private MockCallable innerCallable;
49+
RateLimitingServerStreamingCallable callableToTest;
50+
51+
@Before
52+
public void setup() throws Exception {
53+
innerCallable = new MockCallable();
54+
callableToTest = new RateLimitingServerStreamingCallable(innerCallable);
55+
}
56+
57+
@Test
58+
public void testWithRateLimitInfo() throws Exception {
59+
callableToTest.call(request, responseObserver, context);
60+
61+
Instant earlier = Instant.now().minus(org.threeten.bp.Duration.ofHours(1));
62+
63+
// make sure QPS will be updated
64+
callableToTest.getLastQpsChangeTime().set(earlier);
65+
double oldQps = callableToTest.getCurrentRate();
66+
67+
double factor = 0.8;
68+
69+
RateLimitInfo info =
70+
RateLimitInfo.newBuilder()
71+
.setFactor(factor)
72+
.setPeriod(Duration.newBuilder().setSeconds(10).build())
73+
.build();
74+
75+
MutateRowsResponse response = MutateRowsResponse.newBuilder().setRateLimitInfo(info).build();
76+
77+
innerCallable.getObserver().onResponse(response);
78+
79+
// Give the thread sometime to update the QPS
80+
Thread.sleep(100);
81+
double newQps = callableToTest.getCurrentRate();
82+
83+
assertThat(newQps).isWithin(0.1).of(oldQps * factor);
84+
85+
innerCallable.getObserver().onComplete();
86+
}
87+
88+
@Test
89+
public void testNoUpdateWithinPeriod() throws Exception {
90+
callableToTest.call(request, responseObserver, context);
91+
92+
Instant now = Instant.now();
93+
// make sure QPS will not be updated
94+
callableToTest.getLastQpsChangeTime().set(now);
95+
double oldQps = callableToTest.getCurrentRate();
96+
97+
double factor = 0.3;
98+
99+
RateLimitInfo info =
100+
RateLimitInfo.newBuilder()
101+
.setFactor(factor)
102+
.setPeriod(Duration.newBuilder().setSeconds(600).build())
103+
.build();
104+
105+
MutateRowsResponse response = MutateRowsResponse.newBuilder().setRateLimitInfo(info).build();
106+
107+
innerCallable.getObserver().onResponse(response);
108+
109+
// Give the thread sometime to update the QPS
110+
Thread.sleep(100);
111+
double newQps = callableToTest.getCurrentRate();
112+
113+
assertThat(newQps).isEqualTo(oldQps);
114+
115+
innerCallable.getObserver().onComplete();
116+
}
117+
118+
@Test
119+
public void testErrorInfoLowerQPS() throws Exception {
120+
callableToTest.call(request, responseObserver, context);
121+
122+
Instant earlier = Instant.now().minus(org.threeten.bp.Duration.ofHours(1));
123+
124+
// make sure QPS will be updated
125+
callableToTest.getLastQpsChangeTime().set(earlier);
126+
double oldQps = callableToTest.getCurrentRate();
127+
128+
innerCallable
129+
.getObserver()
130+
.onError(
131+
new DeadlineExceededException(
132+
new Throwable(), new FakeStatusCode(StatusCode.Code.DEADLINE_EXCEEDED), false));
133+
134+
// Give the thread sometime to update the QPS
135+
Thread.sleep(100);
136+
double newQps = callableToTest.getCurrentRate();
137+
138+
assertThat(newQps).isWithin(0.1).of(oldQps * RateLimitingServerStreamingCallable.MIN_FACTOR);
139+
}
140+
141+
private static class MockResponseObserver implements ResponseObserver<MutateRowsResponse> {
142+
143+
private ResponseObserver<MutateRowsResponse> observer;
144+
145+
MockResponseObserver(ResponseObserver<MutateRowsResponse> responseObserver) {
146+
this.observer = responseObserver;
147+
}
148+
149+
@Override
150+
public void onStart(StreamController streamController) {
151+
observer.onStart(streamController);
152+
}
153+
154+
@Override
155+
public void onResponse(MutateRowsResponse o) {
156+
observer.onResponse(o);
157+
}
158+
159+
@Override
160+
public void onError(Throwable throwable) {
161+
observer.onError(throwable);
162+
}
163+
164+
@Override
165+
public void onComplete() {
166+
observer.onComplete();
167+
}
168+
}
169+
170+
private static class MockCallable
171+
extends ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> {
172+
173+
private ResponseObserver<MutateRowsResponse> observer;
174+
175+
@Override
176+
public void call(
177+
MutateRowsRequest mutateRowsRequest,
178+
ResponseObserver<MutateRowsResponse> responseObserver,
179+
ApiCallContext apiCallContext) {
180+
observer = new MockResponseObserver(responseObserver);
181+
}
182+
183+
ResponseObserver<MutateRowsResponse> getObserver() {
184+
return observer;
185+
}
186+
}
187+
}

0 commit comments

Comments
 (0)