This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8610459f7b8 [Gemini] Port client-side throttling to the Java SDK
(#39021)
8610459f7b8 is described below
commit 8610459f7b8c33bb852e97c38236dadf67a2baa3
Author: Jack McCluskey <[email protected]>
AuthorDate: Wed Jul 1 14:57:50 2026 -0400
[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>
---
.../components/throttling/AdaptiveThrottler.java | 105 +++++++++++++++++++
.../components/throttling/ReactiveThrottler.java | 78 ++++++++++++++
.../throttling/AdaptiveThrottlerTest.java | 107 +++++++++++++++++++
sdks/java/io/google-cloud-platform/build.gradle | 1 +
.../sdk/io/gcp/datastore/AdaptiveThrottler.java | 111 --------------------
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 1 +
.../io/gcp/datastore/AdaptiveThrottlerTest.java | 114 ---------------------
7 files changed, 292 insertions(+), 225 deletions(-)
diff --git
a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java
b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java
new file mode 100644
index 00000000000..b2c94a5de35
--- /dev/null
+++
b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.components.throttling;
+
+import java.util.Random;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.util.MovingFunction;
+
+/**
+ * Implements adaptive throttling.
+ *
+ * <p>See
+ *
https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg
+ * for a full discussion of the use case and algorithm applied.
+ */
+public class AdaptiveThrottler {
+
+ // The target minimum number of requests per samplePeriodMs, even if no
+ // requests succeed. Must be greater than 0, else we could throttle to zero.
+ // Because every decision is probabilistic, there is no guarantee that the
+ // request rate in any given interval will not be zero. (This is the +1 from
+ // the formula in
+ // https://landing.google.com/sre/book/chapters/handling-overload.html )
+ public static final int MIN_REQUESTS = 1;
+
+ private final MovingFunction allRequests;
+ private final MovingFunction successfulRequests;
+ private final double overloadRatio;
+ private final Random random;
+
+ /**
+ * Initializes AdaptiveThrottler.
+ *
+ * @param samplePeriodMs length of history to consider, in ms, to set
throttling.
+ * @param sampleUpdateMs granularity of time buckets that we store data in,
in ms.
+ * @param overloadRatio the target ratio between requests sent and
successful requests.
+ */
+ public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double
overloadRatio) {
+ this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random());
+ }
+
+ // visible for testing
+ AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double
overloadRatio, Random random) {
+ if (overloadRatio <= 1.0) {
+ throw new IllegalArgumentException("overloadRatio must be greater than
1.0");
+ }
+ this.allRequests = new MovingFunction(samplePeriodMs, sampleUpdateMs, 1,
1, Sum.ofLongs());
+ this.successfulRequests =
+ new MovingFunction(samplePeriodMs, sampleUpdateMs, 1, 1,
Sum.ofLongs());
+ this.overloadRatio = overloadRatio;
+ this.random = random;
+ }
+
+ protected double throttlingProbability(long nowMsSinceEpoch) {
+ long allReqs = allRequests.get(nowMsSinceEpoch);
+ if (!allRequests.isSignificant()) {
+ return 0.0;
+ }
+ long successfulReqs = successfulRequests.get(nowMsSinceEpoch);
+ double prob = (allReqs - overloadRatio * successfulReqs) / (allReqs +
MIN_REQUESTS);
+ return Math.max(0.0, prob);
+ }
+
+ /**
+ * Determines whether one RPC attempt should be throttled.
+ *
+ * <p>This should be called once each time the caller intends to send an
RPC; if it returns true,
+ * drop or delay that request (calling this function again after the delay).
+ *
+ * @param nowMsSinceEpoch time in ms since the epoch
+ * @return true if the caller should throttle or delay the request.
+ */
+ public synchronized boolean throttleRequest(long nowMsSinceEpoch) {
+ double prob = throttlingProbability(nowMsSinceEpoch);
+ allRequests.add(nowMsSinceEpoch, 1);
+ return random.nextDouble() < prob;
+ }
+
+ /**
+ * Notifies the throttler of a successful request.
+ *
+ * <p>Must be called once for each request (for which throttleRequest was
previously called) that
+ * succeeded.
+ *
+ * @param nowMsSinceEpoch time in ms since the epoch
+ */
+ public synchronized void successfulRequest(long nowMsSinceEpoch) {
+ successfulRequests.add(nowMsSinceEpoch, 1);
+ }
+}
diff --git
a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java
b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java
new file mode 100644
index 00000000000..2c828f9c918
--- /dev/null
+++
b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.components.throttling;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A wrapper around the AdaptiveThrottler that also handles logging and
signaling throttling to the
+ * SDK harness using the provided namespace.
+ *
+ * <p>For usage, instantiate one instance of a ReactiveThrottler class for a
PTransform. When making
+ * remote calls to a service, preface that call with the throttle() method to
potentially
+ * pre-emptively throttle the request. This will throttle future calls based
on the failure rate of
+ * preceding calls, with higher failure rates leading to longer periods of
throttling to allow
+ * system recovery. capture the timestamp of the attempted request, then
execute the request code.
+ * On a success, call successfulRequest(timestamp) to report the success to
the throttler.
+ */
+public class ReactiveThrottler extends AdaptiveThrottler {
+ private static final Logger LOG =
LoggerFactory.getLogger(ReactiveThrottler.class);
+ private static final long SECONDS_TO_MILLISECONDS = 1000L;
+
+ private final ThrottlingSignaler throttlingSignaler;
+ private final int throttleDelaySecs;
+
+ /**
+ * Initializes the ReactiveThrottler.
+ *
+ * @param samplePeriodMs length of history to consider, in ms, to set
throttling.
+ * @param sampleUpdateMs granularity of time buckets that we store data in,
in ms.
+ * @param overloadRatio the target ratio between requests sent and
successful requests.
+ * @param namespace the namespace to use for logging and signaling
throttling is occurring.
+ * @param throttleDelaySecs the amount of time in seconds to wait after
preemptively throttled
+ * requests.
+ */
+ public ReactiveThrottler(
+ long samplePeriodMs,
+ long sampleUpdateMs,
+ double overloadRatio,
+ String namespace,
+ int throttleDelaySecs) {
+ super(samplePeriodMs, sampleUpdateMs, overloadRatio);
+ if (throttleDelaySecs <= 0) {
+ throw new IllegalArgumentException("throttleDelaySecs must be greater
than 0");
+ }
+ this.throttlingSignaler = new ThrottlingSignaler(namespace);
+ this.throttleDelaySecs = throttleDelaySecs;
+ }
+
+ /**
+ * Stops request code from advancing while the underlying AdaptiveThrottler
is signaling to
+ * preemptively throttle the request. Automatically handles logging the
throttling and signaling
+ * to the SDK harness that the request is being throttled. This should be
called in any context
+ * where a call to a remote service is being contacted prior to the call
being performed.
+ */
+ public void throttle() throws InterruptedException {
+ if (throttleRequest(System.currentTimeMillis())) {
+ LOG.debug("Delaying request for {} seconds due to previous failures",
throttleDelaySecs);
+ Thread.sleep(throttleDelaySecs * SECONDS_TO_MILLISECONDS);
+ throttlingSignaler.signalThrottling(throttleDelaySecs *
SECONDS_TO_MILLISECONDS);
+ }
+ }
+}
diff --git
a/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottlerTest.java
b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottlerTest.java
new file mode 100644
index 00000000000..3e798c7a0d5
--- /dev/null
+++
b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottlerTest.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.components.throttling;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Random;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class AdaptiveThrottlerTest {
+
+ private static final long START_TIME = 1500000000000L;
+ private static final long SAMPLE_PERIOD = 60000L;
+ private static final long BUCKET = 1000L;
+ private static final double OVERLOAD_RATIO = 2.0;
+
+ private AdaptiveThrottler throttler;
+
+ @Before
+ public void setUp() {
+ throttler = new AdaptiveThrottler(SAMPLE_PERIOD, BUCKET, OVERLOAD_RATIO);
+ }
+
+ @Test
+ public void testNoInitialThrottling() {
+ assertEquals(0.0, throttler.throttlingProbability(START_TIME), 0.0);
+ }
+
+ @Test
+ public void testNoThrottlingIfNoErrors() {
+ for (long t = START_TIME; t < START_TIME + 20; t++) {
+ assertFalse(throttler.throttleRequest(t));
+ throttler.successfulRequest(t);
+ }
+ assertEquals(0.0, throttler.throttlingProbability(START_TIME + 20), 0.0);
+ }
+
+ @Test
+ public void testNoThrottlingAfterErrorsExpire() {
+ for (long t = START_TIME; t < START_TIME + SAMPLE_PERIOD; t += 100) {
+ throttler.throttleRequest(t);
+ // No successful request
+ }
+ assertTrue(throttler.throttlingProbability(START_TIME + SAMPLE_PERIOD) >
0);
+
+ for (long t = START_TIME + SAMPLE_PERIOD; t < START_TIME + SAMPLE_PERIOD *
2; t += 100) {
+ throttler.throttleRequest(t);
+ throttler.successfulRequest(t);
+ }
+
+ assertEquals(0.0, throttler.throttlingProbability(START_TIME +
SAMPLE_PERIOD * 2), 0.0);
+ }
+
+ @Test
+ public void testThrottlingAfterErrors() {
+ // Inject a mocked Random
+ throttler = new AdaptiveThrottler(SAMPLE_PERIOD, BUCKET, OVERLOAD_RATIO,
new MockRandom());
+
+ for (long t = START_TIME; t < START_TIME + 20; t++) {
+ boolean throttled = throttler.throttleRequest(t);
+ // 1/3rd of requests succeeding.
+ if (t % 3 == 1) {
+ throttler.successfulRequest(t);
+ }
+
+ if (t > START_TIME + 10) {
+ // Roughly 1/3rd succeeding, 1/3rd failing, 1/3rd throttled.
+ assertEquals(0.33, throttler.throttlingProbability(t), 0.1);
+ // Given mocked random, expects 10..13 throttled, 14+ unthrottled
+ assertEquals(t < START_TIME + 14, throttled);
+ }
+ }
+ }
+
+ private static class MockRandom extends Random {
+ private int callCount = 0;
+
+ @Override
+ public double nextDouble() {
+ // Return 0.0, 0.1, ..., 0.9, 0.0, 0.1 ...
+ double val = (callCount % 10) / 10.0;
+ callCount++;
+ return val;
+ }
+ }
+}
diff --git a/sdks/java/io/google-cloud-platform/build.gradle
b/sdks/java/io/google-cloud-platform/build.gradle
index 2686297c000..c8cc9958e2f 100644
--- a/sdks/java/io/google-cloud-platform/build.gradle
+++ b/sdks/java/io/google-cloud-platform/build.gradle
@@ -41,6 +41,7 @@ dependencies {
permitUnusedDeclared project(":sdks:java:expansion-service") // BEAM-11761
implementation project(":sdks:java:extensions:google-cloud-platform-core")
implementation project(":sdks:java:extensions:protobuf")
+ implementation project(":sdks:java:io:components")
implementation project(":sdks:java:extensions:arrow")
implementation library.java.avro
implementation library.java.bigdataoss_util
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
deleted file mode 100644
index 15c4e678bcb..00000000000
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.datastore;
-
-import java.util.Random;
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.util.MovingFunction;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-
-/**
- * An implementation of client-side adaptive throttling. See
- *
https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg
- * for a full discussion of the use case and algorithm applied.
- */
-class AdaptiveThrottler {
- private final MovingFunction successfulRequests;
- private final MovingFunction allRequests;
-
- /**
- * The target ratio between requests sent and successful requests. This is
"K" in the formula in
- * https://landing.google.com/sre/book/chapters/handling-overload.html
- */
- private final double overloadRatio;
-
- /**
- * The target minimum number of requests per samplePeriodMs, even if no
requests succeed. Must be
- * greater than 0, else we could throttle to zero. Because every decision is
probabilistic, there
- * is no guarantee that the request rate in any given interval will not be
zero. (This is the +1
- * from the formula in
https://landing.google.com/sre/book/chapters/handling-overload.html
- */
- private static final double MIN_REQUESTS = 1;
-
- private final Random random;
-
- /**
- * @param samplePeriodMs the time window to keep of request history to
inform throttling
- * decisions.
- * @param sampleUpdateMs the length of buckets within this time window.
- * @param overloadRatio the target ratio between requests sent and
successful requests. You should
- * always set this to more than 1, otherwise the client would never try
to send more requests
- * than succeeded in the past - so it could never recover from temporary
setbacks.
- */
- public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double
overloadRatio) {
- this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random());
- }
-
- @VisibleForTesting
- AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double
overloadRatio, Random random) {
- allRequests =
- new MovingFunction(
- samplePeriodMs,
- sampleUpdateMs,
- 1 /* numSignificantBuckets */,
- 1 /* numSignificantSamples */,
- Sum.ofLongs());
- successfulRequests =
- new MovingFunction(
- samplePeriodMs,
- sampleUpdateMs,
- 1 /* numSignificantBuckets */,
- 1 /* numSignificantSamples */,
- Sum.ofLongs());
- this.overloadRatio = overloadRatio;
- this.random = random;
- }
-
- @VisibleForTesting
- double throttlingProbability(long nowMsSinceEpoch) {
- if (!allRequests.isSignificant()) {
- return 0;
- }
- long allRequestsNow = allRequests.get(nowMsSinceEpoch);
- long successfulRequestsNow = successfulRequests.get(nowMsSinceEpoch);
- return Math.max(
- 0,
- (allRequestsNow - overloadRatio * successfulRequestsNow) /
(allRequestsNow + MIN_REQUESTS));
- }
-
- /**
- * Call this before sending a request to the remote service; if this returns
true, drop the
- * request (treating it as a failure or trying it again at a later time).
- */
- public boolean throttleRequest(long nowMsSinceEpoch) {
- double delayProbability = throttlingProbability(nowMsSinceEpoch);
- // Note that we increment the count of all requests here, even if we
return true - so even if we
- // tell the client not to send a request at all, it still counts as a
failed request.
- allRequests.add(nowMsSinceEpoch, 1);
-
- return (random.nextDouble() < delayProbability);
- }
-
- /** Call this after {@link throttleRequest} if your request was successful.
*/
- public void successfulRequest(long nowMsSinceEpoch) {
- successfulRequests.add(nowMsSinceEpoch, 1);
- }
-}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
index c9507475648..6966de5b484 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
@@ -77,6 +77,7 @@ import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer;
+import org.apache.beam.sdk.io.components.throttling.AdaptiveThrottler;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Distribution;
import org.apache.beam.sdk.metrics.Metrics;
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
deleted file mode 100644
index 6db8af3906c..00000000000
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
+++ /dev/null
@@ -1,114 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.datastore;
-
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.closeTo;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.junit.Assert.assertFalse;
-
-import java.util.Random;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
-
-/** Tests for {@link AdaptiveThrottler}. */
-@RunWith(JUnit4.class)
-public class AdaptiveThrottlerTest {
-
- static final long START_TIME_MS = 0;
- static final long SAMPLE_PERIOD_MS = 60000;
- static final long SAMPLE_BUCKET_MS = 1000;
- static final double OVERLOAD_RATIO = 2;
-
- /** Returns a throttler configured with the standard parameters above. */
- AdaptiveThrottler getThrottler() {
- return new AdaptiveThrottler(SAMPLE_PERIOD_MS, SAMPLE_BUCKET_MS,
OVERLOAD_RATIO);
- }
-
- @Test
- public void testNoInitialThrottling() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- assertThat(throttler.throttlingProbability(START_TIME_MS), equalTo(0.0));
- assertThat(
- "first request is not throttled",
throttler.throttleRequest(START_TIME_MS), equalTo(false));
- }
-
- @Test
- public void testNoThrottlingIfNoErrors() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- long t = START_TIME_MS;
- for (; t < START_TIME_MS + 20; t++) {
- assertFalse(throttler.throttleRequest(t));
- throttler.successfulRequest(t);
- }
- assertThat(throttler.throttlingProbability(t), equalTo(0.0));
- }
-
- @Test
- public void testNoThrottlingAfterErrorsExpire() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- long t = START_TIME_MS;
- for (; t < START_TIME_MS + SAMPLE_PERIOD_MS; t++) {
- throttler.throttleRequest(t);
- // and no successfulRequest.
- }
- assertThat(
- "check that we set up a non-zero probability of throttling",
- throttler.throttlingProbability(t),
- greaterThan(0.0));
- for (; t < START_TIME_MS + 2 * SAMPLE_PERIOD_MS; t++) {
- throttler.throttleRequest(t);
- throttler.successfulRequest(t);
- }
- assertThat(throttler.throttlingProbability(t), equalTo(0.0));
- }
-
- @Test
- public void testThrottlingAfterErrors() throws Exception {
- Random mockRandom = Mockito.mock(Random.class);
- Mockito.when(mockRandom.nextDouble())
- .thenReturn(
- 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.0, 0.1, 0.2,
0.3, 0.4, 0.5, 0.6,
- 0.7, 0.8, 0.9);
- AdaptiveThrottler throttler =
- new AdaptiveThrottler(SAMPLE_PERIOD_MS, SAMPLE_BUCKET_MS,
OVERLOAD_RATIO, mockRandom);
- for (int i = 0; i < 20; i++) {
- boolean throttled = throttler.throttleRequest(START_TIME_MS + i);
- // 1/3rd of requests succeeding.
- if (i % 3 == 1) {
- throttler.successfulRequest(START_TIME_MS + i);
- }
-
- // Once we have some history in place, check what throttling happens.
- if (i >= 10) {
- // Expect 1/3rd of requests to be throttled. (So 1/3rd throttled,
1/3rd succeeding, 1/3rd
- // tried and failing).
- assertThat(
- String.format("for i=%d", i),
- throttler.throttlingProbability(START_TIME_MS + i),
- closeTo(0.33, /*error=*/ 0.1));
- // Requests 10..13 should be throttled, 14..19 not throttled given the
mocked random numbers
- // that we fed to throttler.
- assertThat(String.format("for i=%d", i), throttled, equalTo(i < 14));
- }
- }
- }
-}
