Skip to content

Commit 8610459

Browse files
jrmccluskeygemini-code-assist[bot]
jrmccluskey
and
gemini-code-assist[bot]
authored
[Gemini] Port client-side throttling to the Java SDK (#39021)
* [Gemini] Port client-side throttling to the Java SDK * Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * code suggestions * package-info * cleanup * consolidate adaptive throttler impls * Remove empty util directory --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 00823b2 commit 8610459

7 files changed

Lines changed: 292 additions & 225 deletions

File tree

‎sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.components.throttling;
19+
20+
import java.util.Random;
21+
import org.apache.beam.sdk.transforms.Sum;
22+
import org.apache.beam.sdk.util.MovingFunction;
23+
24+
/**
25+
* Implements adaptive throttling.
26+
*
27+
* <p>See
28+
* https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg
29+
* for a full discussion of the use case and algorithm applied.
30+
*/
31+
public class AdaptiveThrottler {
32+
33+
// The target minimum number of requests per samplePeriodMs, even if no
34+
// requests succeed. Must be greater than 0, else we could throttle to zero.
35+
// Because every decision is probabilistic, there is no guarantee that the
36+
// request rate in any given interval will not be zero. (This is the +1 from
37+
// the formula in
38+
// https://landing.google.com/sre/book/chapters/handling-overload.html )
39+
public static final int MIN_REQUESTS = 1;
40+
41+
private final MovingFunction allRequests;
42+
private final MovingFunction successfulRequests;
43+
private final double overloadRatio;
44+
private final Random random;
45+
46+
/**
47+
* Initializes AdaptiveThrottler.
48+
*
49+
* @param samplePeriodMs length of history to consider, in ms, to set throttling.
50+
* @param sampleUpdateMs granularity of time buckets that we store data in, in ms.
51+
* @param overloadRatio the target ratio between requests sent and successful requests.
52+
*/
53+
public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double overloadRatio) {
54+
this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random());
55+
}
56+
57+
// visible for testing
58+
AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double overloadRatio, Random random) {
59+
if (overloadRatio <= 1.0) {
60+
throw new IllegalArgumentException("overloadRatio must be greater than 1.0");
61+
}
62+
this.allRequests = new MovingFunction(samplePeriodMs, sampleUpdateMs, 1, 1, Sum.ofLongs());
63+
this.successfulRequests =
64+
new MovingFunction(samplePeriodMs, sampleUpdateMs, 1, 1, Sum.ofLongs());
65+
this.overloadRatio = overloadRatio;
66+
this.random = random;
67+
}
68+
69+
protected double throttlingProbability(long nowMsSinceEpoch) {
70+
long allReqs = allRequests.get(nowMsSinceEpoch);
71+
if (!allRequests.isSignificant()) {
72+
return 0.0;
73+
}
74+
long successfulReqs = successfulRequests.get(nowMsSinceEpoch);
75+
double prob = (allReqs - overloadRatio * successfulReqs) / (allReqs + MIN_REQUESTS);
76+
return Math.max(0.0, prob);
77+
}
78+
79+
/**
80+
* Determines whether one RPC attempt should be throttled.
81+
*
82+
* <p>This should be called once each time the caller intends to send an RPC; if it returns true,
83+
* drop or delay that request (calling this function again after the delay).
84+
*
85+
* @param nowMsSinceEpoch time in ms since the epoch
86+
* @return true if the caller should throttle or delay the request.
87+
*/
88+
public synchronized boolean throttleRequest(long nowMsSinceEpoch) {
89+
double prob = throttlingProbability(nowMsSinceEpoch);
90+
allRequests.add(nowMsSinceEpoch, 1);
91+
return random.nextDouble() < prob;
92+
}
93+
94+
/**
95+
* Notifies the throttler of a successful request.
96+
*
97+
* <p>Must be called once for each request (for which throttleRequest was previously called) that
98+
* succeeded.
99+
*
100+
* @param nowMsSinceEpoch time in ms since the epoch
101+
*/
102+
public synchronized void successfulRequest(long nowMsSinceEpoch) {
103+
successfulRequests.add(nowMsSinceEpoch, 1);
104+
}
105+
}

‎sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.components.throttling;
19+
20+
import org.slf4j.Logger;
21+
import org.slf4j.LoggerFactory;
22+
23+
/**
24+
* A wrapper around the AdaptiveThrottler that also handles logging and signaling throttling to the
25+
* SDK harness using the provided namespace.
26+
*
27+
* <p>For usage, instantiate one instance of a ReactiveThrottler class for a PTransform. When making
28+
* remote calls to a service, preface that call with the throttle() method to potentially
29+
* pre-emptively throttle the request. This will throttle future calls based on the failure rate of
30+
* preceding calls, with higher failure rates leading to longer periods of throttling to allow
31+
* system recovery. capture the timestamp of the attempted request, then execute the request code.
32+
* On a success, call successfulRequest(timestamp) to report the success to the throttler.
33+
*/
34+
public class ReactiveThrottler extends AdaptiveThrottler {
35+
private static final Logger LOG = LoggerFactory.getLogger(ReactiveThrottler.class);
36+
private static final long SECONDS_TO_MILLISECONDS = 1000L;
37+
38+
private final ThrottlingSignaler throttlingSignaler;
39+
private final int throttleDelaySecs;
40+
41+
/**
42+
* Initializes the ReactiveThrottler.
43+
*
44+
* @param samplePeriodMs length of history to consider, in ms, to set throttling.
45+
* @param sampleUpdateMs granularity of time buckets that we store data in, in ms.
46+
* @param overloadRatio the target ratio between requests sent and successful requests.
47+
* @param namespace the namespace to use for logging and signaling throttling is occurring.
48+
* @param throttleDelaySecs the amount of time in seconds to wait after preemptively throttled
49+
* requests.
50+
*/
51+
public ReactiveThrottler(
52+
long samplePeriodMs,
53+
long sampleUpdateMs,
54+
double overloadRatio,
55+
String namespace,
56+
int throttleDelaySecs) {
57+
super(samplePeriodMs, sampleUpdateMs, overloadRatio);
58+
if (throttleDelaySecs <= 0) {
59+
throw new IllegalArgumentException("throttleDelaySecs must be greater than 0");
60+
}
61+
this.throttlingSignaler = new ThrottlingSignaler(namespace);
62+
this.throttleDelaySecs = throttleDelaySecs;
63+
}
64+
65+
/**
66+
* Stops request code from advancing while the underlying AdaptiveThrottler is signaling to
67+
* preemptively throttle the request. Automatically handles logging the throttling and signaling
68+
* to the SDK harness that the request is being throttled. This should be called in any context
69+
* where a call to a remote service is being contacted prior to the call being performed.
70+
*/
71+
public void throttle() throws InterruptedException {
72+
if (throttleRequest(System.currentTimeMillis())) {
73+
LOG.debug("Delaying request for {} seconds due to previous failures", throttleDelaySecs);
74+
Thread.sleep(throttleDelaySecs * SECONDS_TO_MILLISECONDS);
75+
throttlingSignaler.signalThrottling(throttleDelaySecs * SECONDS_TO_MILLISECONDS);
76+
}
77+
}
78+
}

‎sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottlerTest.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.components.throttling;
19+
20+
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertFalse;
22+
import static org.junit.Assert.assertTrue;
23+
24+
import java.util.Random;
25+
import org.junit.Before;
26+
import org.junit.Test;
27+
import org.junit.runner.RunWith;
28+
import org.junit.runners.JUnit4;
29+
30+
@RunWith(JUnit4.class)
31+
public class AdaptiveThrottlerTest {
32+
33+
private static final long START_TIME = 1500000000000L;
34+
private static final long SAMPLE_PERIOD = 60000L;
35+
private static final long BUCKET = 1000L;
36+
private static final double OVERLOAD_RATIO = 2.0;
37+
38+
private AdaptiveThrottler throttler;
39+
40+
@Before
41+
public void setUp() {
42+
throttler = new AdaptiveThrottler(SAMPLE_PERIOD, BUCKET, OVERLOAD_RATIO);
43+
}
44+
45+
@Test
46+
public void testNoInitialThrottling() {
47+
assertEquals(0.0, throttler.throttlingProbability(START_TIME), 0.0);
48+
}
49+
50+
@Test
51+
public void testNoThrottlingIfNoErrors() {
52+
for (long t = START_TIME; t < START_TIME + 20; t++) {
53+
assertFalse(throttler.throttleRequest(t));
54+
throttler.successfulRequest(t);
55+
}
56+
assertEquals(0.0, throttler.throttlingProbability(START_TIME + 20), 0.0);
57+
}
58+
59+
@Test
60+
public void testNoThrottlingAfterErrorsExpire() {
61+
for (long t = START_TIME; t < START_TIME + SAMPLE_PERIOD; t += 100) {
62+
throttler.throttleRequest(t);
63+
// No successful request
64+
}
65+
assertTrue(throttler.throttlingProbability(START_TIME + SAMPLE_PERIOD) > 0);
66+
67+
for (long t = START_TIME + SAMPLE_PERIOD; t < START_TIME + SAMPLE_PERIOD * 2; t += 100) {
68+
throttler.throttleRequest(t);
69+
throttler.successfulRequest(t);
70+
}
71+
72+
assertEquals(0.0, throttler.throttlingProbability(START_TIME + SAMPLE_PERIOD * 2), 0.0);
73+
}
74+
75+
@Test
76+
public void testThrottlingAfterErrors() {
77+
// Inject a mocked Random
78+
throttler = new AdaptiveThrottler(SAMPLE_PERIOD, BUCKET, OVERLOAD_RATIO, new MockRandom());
79+
80+
for (long t = START_TIME; t < START_TIME + 20; t++) {
81+
boolean throttled = throttler.throttleRequest(t);
82+
// 1/3rd of requests succeeding.
83+
if (t % 3 == 1) {
84+
throttler.successfulRequest(t);
85+
}
86+
87+
if (t > START_TIME + 10) {
88+
// Roughly 1/3rd succeeding, 1/3rd failing, 1/3rd throttled.
89+
assertEquals(0.33, throttler.throttlingProbability(t), 0.1);
90+
// Given mocked random, expects 10..13 throttled, 14+ unthrottled
91+
assertEquals(t < START_TIME + 14, throttled);
92+
}
93+
}
94+
}
95+
96+
private static class MockRandom extends Random {
97+
private int callCount = 0;
98+
99+
@Override
100+
public double nextDouble() {
101+
// Return 0.0, 0.1, ..., 0.9, 0.0, 0.1 ...
102+
double val = (callCount % 10) / 10.0;
103+
callCount++;
104+
return val;
105+
}
106+
}
107+
}

‎sdks/java/io/google-cloud-platform/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ dependencies {
4141
permitUnusedDeclared project(":sdks:java:expansion-service") // BEAM-11761
4242
implementation project(":sdks:java:extensions:google-cloud-platform-core")
4343
implementation project(":sdks:java:extensions:protobuf")
44+
implementation project(":sdks:java:io:components")
4445
implementation project(":sdks:java:extensions:arrow")
4546
implementation library.java.avro
4647
implementation library.java.bigdataoss_util

0 commit comments

Comments
 (0)