From 97671239a98c1f8bbfc809aaaa78036706cbc2ea Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 18 Jun 2026 11:50:57 -0400 Subject: [PATCH 1/7] [Gemini] Port client-side throttling to the Java SDK --- .../throttling/AdaptiveThrottler.java | 100 ++++++++++++++++ .../throttling/ReactiveThrottler.java | 71 ++++++++++++ .../sdk/io/components/util/MovingSum.java | 104 +++++++++++++++++ .../throttling/AdaptiveThrottlerTest.java | 107 ++++++++++++++++++ .../sdk/io/components/util/MovingSumTest.java | 57 ++++++++++ 5 files changed, 439 insertions(+) create mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java create mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java create mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java create mode 100644 sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottlerTest.java create mode 100644 sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java 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 000000000000..dff16a313058 --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/AdaptiveThrottler.java @@ -0,0 +1,100 @@ +/* + * 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.io.components.util.MovingSum; + +/** + * Implements 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. + */ +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 MovingSum allRequests; + private final MovingSum successfulRequests; + private final double overloadRatio; + private final Random random; + + /** + * Initializes AdaptiveThrottler. + * + * @param windowMs length of history to consider, in ms, to set throttling. + * @param bucketMs 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 windowMs, long bucketMs, double overloadRatio) { + this(windowMs, bucketMs, overloadRatio, new Random()); + } + + // visible for testing + AdaptiveThrottler(long windowMs, long bucketMs, double overloadRatio, Random random) { + this.allRequests = new MovingSum(windowMs, bucketMs); + this.successfulRequests = new MovingSum(windowMs, bucketMs); + this.overloadRatio = overloadRatio; + this.random = random; + } + + protected double throttlingProbability(long now) { + if (!allRequests.hasData(now)) { + return 0.0; + } + long allReqs = allRequests.sum(now); + long successfulReqs = successfulRequests.sum(now); + double prob = (allReqs - overloadRatio * successfulReqs) / (allReqs + MIN_REQUESTS); + return Math.max(0.0, prob); + } + + /** + * Determines whether one RPC attempt should be throttled. + * + *

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 now time in ms since the epoch + * @return true if the caller should throttle or delay the request. + */ + public boolean throttleRequest(long now) { + double prob = throttlingProbability(now); + allRequests.add(now, 1); + return random.nextDouble() < prob; + } + + /** + * Notifies the throttler of a successful request. + * + *

Must be called once for each request (for which throttleRequest was previously called) that + * succeeded. + * + * @param now time in ms since the epoch + */ + public void successfulRequest(long now) { + successfulRequests.add(now, 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 000000000000..c8fb681d059f --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/throttling/ReactiveThrottler.java @@ -0,0 +1,71 @@ +/* + * 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. + * + *

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 windowMs length of history to consider, in ms, to set throttling. + * @param bucketMs 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 windowMs, long bucketMs, double overloadRatio, String namespace, int throttleDelaySecs) { + super(windowMs, bucketMs, overloadRatio); + 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 { + while (throttleRequest(System.currentTimeMillis())) { + LOG.info("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/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java new file mode 100644 index 000000000000..63d7eaabe39f --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java @@ -0,0 +1,104 @@ +/* + * 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.util; + +import java.util.Arrays; + +/** + * Class that keeps track of a rolling window sum. + * + *

For use in tracking recent performance of the connector. + * + *

Intended to be similar to {@link org.apache.beam.sdk.util.MovingFunction}, but for convenience + * we expose the count of entries as well so this doubles as a moving average tracker. + */ +public class MovingSum { + private final int numBuckets; + private final long bucketMs; + + private int currentIndex; + private long currentMsSinceEpoch; + private final long[] sums; + private final long[] counts; + + public MovingSum(long windowMs, long bucketMs) { + if (windowMs < bucketMs || bucketMs <= 0) { + throw new IllegalArgumentException("windowMs >= bucketMs > 0 please"); + } + this.numBuckets = (int) Math.ceil((double) windowMs / bucketMs); + this.bucketMs = bucketMs; + this.sums = new long[this.numBuckets]; + this.counts = new long[this.numBuckets]; + this.currentIndex = 0; + this.currentMsSinceEpoch = 0; + Arrays.fill(this.sums, 0L); + Arrays.fill(this.counts, 0L); + } + + private void reset(long now) { + this.currentIndex = 0; + this.currentMsSinceEpoch = (now / bucketMs) * bucketMs; + Arrays.fill(sums, 0L); + Arrays.fill(counts, 0L); + } + + private void flush(long now) { + if (now >= (currentMsSinceEpoch + bucketMs * numBuckets)) { + // Time moved forward so far that all currently held data is outside of + // the window. It is faster to simply reset our data. + reset(now); + return; + } + + while (now > currentMsSinceEpoch + bucketMs) { + // Advance time by one bucketMs, setting the new bucket's counts to 0. + currentMsSinceEpoch += bucketMs; + currentIndex = (currentIndex + 1) % numBuckets; + sums[currentIndex] = 0; + counts[currentIndex] = 0; + } + } + + public long sum(long now) { + flush(now); + long total = 0; + for (long s : sums) { + total += s; + } + return total; + } + + public void add(long now, long inc) { + flush(now); + sums[currentIndex] += inc; + counts[currentIndex] += 1; + } + + public long count(long now) { + flush(now); + long total = 0; + for (long c : counts) { + total += c; + } + return total; + } + + public boolean hasData(long now) { + return count(now) > 0; + } +} 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 000000000000..3e798c7a0d55 --- /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/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java new file mode 100644 index 000000000000..4f3718fb0a8c --- /dev/null +++ b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java @@ -0,0 +1,57 @@ +/* + * 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.util; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class MovingSumTest { + + @Test + public void testMovingSumBasic() { + MovingSum movingSum = new MovingSum(10000, 1000); + assertFalse(movingSum.hasData(1000)); + + movingSum.add(1000, 5); + assertTrue(movingSum.hasData(1000)); + assertEquals(5, movingSum.sum(1000)); + assertEquals(1, movingSum.count(1000)); + + movingSum.add(1500, 10); + assertEquals(15, movingSum.sum(1500)); + assertEquals(2, movingSum.count(1500)); + + // Wait 11 seconds (moving completely outside window) + assertEquals(0, movingSum.sum(12000)); + assertEquals(0, movingSum.count(12000)); + assertFalse(movingSum.hasData(12000)); + } + + @Test + public void testInvalidArguments() { + assertThrows(IllegalArgumentException.class, () -> new MovingSum(100, 1000)); + assertThrows(IllegalArgumentException.class, () -> new MovingSum(1000, 0)); + } +} From 2827eed869a1ee6c1a60f46e7a99c0fd93d0a252 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Thu, 18 Jun 2026 12:00:25 -0400 Subject: [PATCH 2/7] Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../beam/sdk/io/components/throttling/AdaptiveThrottler.java | 4 ++-- .../beam/sdk/io/components/throttling/ReactiveThrottler.java | 4 ++-- .../org/apache/beam/sdk/io/components/util/MovingSum.java | 2 +- 3 files changed, 5 insertions(+), 5 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 index dff16a313058..d3a38bbe6e38 100644 --- 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 @@ -80,7 +80,7 @@ protected double throttlingProbability(long now) { * @param now time in ms since the epoch * @return true if the caller should throttle or delay the request. */ - public boolean throttleRequest(long now) { + public synchronized boolean throttleRequest(long now) { double prob = throttlingProbability(now); allRequests.add(now, 1); return random.nextDouble() < prob; @@ -94,7 +94,7 @@ public boolean throttleRequest(long now) { * * @param now time in ms since the epoch */ - public void successfulRequest(long now) { + public synchronized void successfulRequest(long now) { successfulRequests.add(now, 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 index c8fb681d059f..9958207c6d03 100644 --- 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 @@ -62,8 +62,8 @@ public ReactiveThrottler( * where a call to a remote service is being contacted prior to the call being performed. */ public void throttle() throws InterruptedException { - while (throttleRequest(System.currentTimeMillis())) { - LOG.info("Delaying request for {} seconds due to previous failures", throttleDelaySecs); + 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/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java index 63d7eaabe39f..92c741d7d668 100644 --- a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java @@ -65,7 +65,7 @@ private void flush(long now) { return; } - while (now > currentMsSinceEpoch + bucketMs) { + while (now >= currentMsSinceEpoch + bucketMs) { // Advance time by one bucketMs, setting the new bucket's counts to 0. currentMsSinceEpoch += bucketMs; currentIndex = (currentIndex + 1) % numBuckets; From 083e2f3863d3ce6f04ab3ab5e7065aaa30976517 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 18 Jun 2026 12:02:49 -0400 Subject: [PATCH 3/7] code suggestions --- .../sdk/io/components/throttling/AdaptiveThrottler.java | 3 +++ .../sdk/io/components/throttling/ReactiveThrottler.java | 3 +++ .../apache/beam/sdk/io/components/util/MovingSumTest.java | 6 ++++++ 3 files changed, 12 insertions(+) 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 index d3a38bbe6e38..58cda496ff6f 100644 --- 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 @@ -55,6 +55,9 @@ public AdaptiveThrottler(long windowMs, long bucketMs, double overloadRatio) { // visible for testing AdaptiveThrottler(long windowMs, long bucketMs, double overloadRatio, Random random) { + if (overloadRatio <= 1.0) { + throw new IllegalArgumentException("overloadRatio must be greater than 1.0"); + } this.allRequests = new MovingSum(windowMs, bucketMs); this.successfulRequests = new MovingSum(windowMs, bucketMs); this.overloadRatio = overloadRatio; 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 index 9958207c6d03..ae0dea81a094 100644 --- 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 @@ -51,6 +51,9 @@ public class ReactiveThrottler extends AdaptiveThrottler { public ReactiveThrottler( long windowMs, long bucketMs, double overloadRatio, String namespace, int throttleDelaySecs) { super(windowMs, bucketMs, overloadRatio); + if (throttleDelaySecs <= 0) { + throw new IllegalArgumentException("throttleDelaySecs must be greater than 0"); + } this.throttlingSignaler = new ThrottlingSignaler(namespace); this.throttleDelaySecs = throttleDelaySecs; } diff --git a/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java index 4f3718fb0a8c..fd70ba7bf241 100644 --- a/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java +++ b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java @@ -43,6 +43,12 @@ public void testMovingSumBasic() { assertEquals(15, movingSum.sum(1500)); assertEquals(2, movingSum.count(1500)); + // Advance by 2 buckets (from 1000 to 3000) + assertEquals(15, movingSum.sum(3000)); + movingSum.add(3500, 20); + assertEquals(35, movingSum.sum(3500)); + assertEquals(3, movingSum.count(3500)); + // Wait 11 seconds (moving completely outside window) assertEquals(0, movingSum.sum(12000)); assertEquals(0, movingSum.count(12000)); From b6bbc0a05610b95ffa6a7456fc4556702a732f8f Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 18 Jun 2026 13:13:18 -0400 Subject: [PATCH 4/7] package-info --- .../sdk/io/components/util/package-info.java | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java new file mode 100644 index 000000000000..cd625446f74e --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Defines utilities for Beam IO components. */ +package org.apache.beam.sdk.io.components.util; \ No newline at end of file From 7b96ba42619cf78a0df4d48dc4d88281c5d7bc4e Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 18 Jun 2026 13:14:24 -0400 Subject: [PATCH 5/7] cleanup --- .../org/apache/beam/sdk/io/components/util/package-info.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java index cd625446f74e..40c757128789 100644 --- a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java @@ -17,4 +17,4 @@ */ /** Defines utilities for Beam IO components. */ -package org.apache.beam.sdk.io.components.util; \ No newline at end of file +package org.apache.beam.sdk.io.components.util; From 1e779a1f331f0d6a9b366a71e64dcb05e78754f6 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 1 Jul 2026 12:14:41 -0400 Subject: [PATCH 6/7] consolidate adaptive throttler impls --- .../throttling/AdaptiveThrottler.java | 44 +++---- .../throttling/ReactiveThrottler.java | 12 +- .../sdk/io/components/util/MovingSum.java | 104 ---------------- .../sdk/io/components/util/MovingSumTest.java | 63 ---------- .../io/google-cloud-platform/build.gradle | 1 + .../io/gcp/datastore/AdaptiveThrottler.java | 111 ----------------- .../sdk/io/gcp/datastore/DatastoreV1.java | 1 + .../gcp/datastore/AdaptiveThrottlerTest.java | 114 ------------------ 8 files changed, 33 insertions(+), 417 deletions(-) delete mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java delete mode 100644 sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java delete mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java delete mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java 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 index 58cda496ff6f..b2c94a5de359 100644 --- 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 @@ -18,7 +18,8 @@ package org.apache.beam.sdk.io.components.throttling; import java.util.Random; -import org.apache.beam.sdk.io.components.util.MovingSum; +import org.apache.beam.sdk.transforms.Sum; +import org.apache.beam.sdk.util.MovingFunction; /** * Implements adaptive throttling. @@ -37,39 +38,40 @@ public class AdaptiveThrottler { // https://landing.google.com/sre/book/chapters/handling-overload.html ) public static final int MIN_REQUESTS = 1; - private final MovingSum allRequests; - private final MovingSum successfulRequests; + private final MovingFunction allRequests; + private final MovingFunction successfulRequests; private final double overloadRatio; private final Random random; /** * Initializes AdaptiveThrottler. * - * @param windowMs length of history to consider, in ms, to set throttling. - * @param bucketMs granularity of time buckets that we store data in, in ms. + * @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 windowMs, long bucketMs, double overloadRatio) { - this(windowMs, bucketMs, overloadRatio, new Random()); + public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs, double overloadRatio) { + this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random()); } // visible for testing - AdaptiveThrottler(long windowMs, long bucketMs, double overloadRatio, Random random) { + 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 MovingSum(windowMs, bucketMs); - this.successfulRequests = new MovingSum(windowMs, bucketMs); + 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 now) { - if (!allRequests.hasData(now)) { + protected double throttlingProbability(long nowMsSinceEpoch) { + long allReqs = allRequests.get(nowMsSinceEpoch); + if (!allRequests.isSignificant()) { return 0.0; } - long allReqs = allRequests.sum(now); - long successfulReqs = successfulRequests.sum(now); + long successfulReqs = successfulRequests.get(nowMsSinceEpoch); double prob = (allReqs - overloadRatio * successfulReqs) / (allReqs + MIN_REQUESTS); return Math.max(0.0, prob); } @@ -80,12 +82,12 @@ protected double throttlingProbability(long now) { *

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 now time in ms since the epoch + * @param nowMsSinceEpoch time in ms since the epoch * @return true if the caller should throttle or delay the request. */ - public synchronized boolean throttleRequest(long now) { - double prob = throttlingProbability(now); - allRequests.add(now, 1); + public synchronized boolean throttleRequest(long nowMsSinceEpoch) { + double prob = throttlingProbability(nowMsSinceEpoch); + allRequests.add(nowMsSinceEpoch, 1); return random.nextDouble() < prob; } @@ -95,9 +97,9 @@ public synchronized boolean throttleRequest(long now) { *

Must be called once for each request (for which throttleRequest was previously called) that * succeeded. * - * @param now time in ms since the epoch + * @param nowMsSinceEpoch time in ms since the epoch */ - public synchronized void successfulRequest(long now) { - successfulRequests.add(now, 1); + 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 index ae0dea81a094..2c828f9c9187 100644 --- 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 @@ -41,16 +41,20 @@ public class ReactiveThrottler extends AdaptiveThrottler { /** * Initializes the ReactiveThrottler. * - * @param windowMs length of history to consider, in ms, to set throttling. - * @param bucketMs granularity of time buckets that we store data in, in ms. + * @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 windowMs, long bucketMs, double overloadRatio, String namespace, int throttleDelaySecs) { - super(windowMs, bucketMs, overloadRatio); + 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"); } diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java deleted file mode 100644 index 92c741d7d668..000000000000 --- a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/MovingSum.java +++ /dev/null @@ -1,104 +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.components.util; - -import java.util.Arrays; - -/** - * Class that keeps track of a rolling window sum. - * - *

For use in tracking recent performance of the connector. - * - *

Intended to be similar to {@link org.apache.beam.sdk.util.MovingFunction}, but for convenience - * we expose the count of entries as well so this doubles as a moving average tracker. - */ -public class MovingSum { - private final int numBuckets; - private final long bucketMs; - - private int currentIndex; - private long currentMsSinceEpoch; - private final long[] sums; - private final long[] counts; - - public MovingSum(long windowMs, long bucketMs) { - if (windowMs < bucketMs || bucketMs <= 0) { - throw new IllegalArgumentException("windowMs >= bucketMs > 0 please"); - } - this.numBuckets = (int) Math.ceil((double) windowMs / bucketMs); - this.bucketMs = bucketMs; - this.sums = new long[this.numBuckets]; - this.counts = new long[this.numBuckets]; - this.currentIndex = 0; - this.currentMsSinceEpoch = 0; - Arrays.fill(this.sums, 0L); - Arrays.fill(this.counts, 0L); - } - - private void reset(long now) { - this.currentIndex = 0; - this.currentMsSinceEpoch = (now / bucketMs) * bucketMs; - Arrays.fill(sums, 0L); - Arrays.fill(counts, 0L); - } - - private void flush(long now) { - if (now >= (currentMsSinceEpoch + bucketMs * numBuckets)) { - // Time moved forward so far that all currently held data is outside of - // the window. It is faster to simply reset our data. - reset(now); - return; - } - - while (now >= currentMsSinceEpoch + bucketMs) { - // Advance time by one bucketMs, setting the new bucket's counts to 0. - currentMsSinceEpoch += bucketMs; - currentIndex = (currentIndex + 1) % numBuckets; - sums[currentIndex] = 0; - counts[currentIndex] = 0; - } - } - - public long sum(long now) { - flush(now); - long total = 0; - for (long s : sums) { - total += s; - } - return total; - } - - public void add(long now, long inc) { - flush(now); - sums[currentIndex] += inc; - counts[currentIndex] += 1; - } - - public long count(long now) { - flush(now); - long total = 0; - for (long c : counts) { - total += c; - } - return total; - } - - public boolean hasData(long now) { - return count(now) > 0; - } -} diff --git a/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java b/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java deleted file mode 100644 index fd70ba7bf241..000000000000 --- a/sdks/java/io/components/src/test/java/org/apache/beam/sdk/io/components/util/MovingSumTest.java +++ /dev/null @@ -1,63 +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.components.util; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; - -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class MovingSumTest { - - @Test - public void testMovingSumBasic() { - MovingSum movingSum = new MovingSum(10000, 1000); - assertFalse(movingSum.hasData(1000)); - - movingSum.add(1000, 5); - assertTrue(movingSum.hasData(1000)); - assertEquals(5, movingSum.sum(1000)); - assertEquals(1, movingSum.count(1000)); - - movingSum.add(1500, 10); - assertEquals(15, movingSum.sum(1500)); - assertEquals(2, movingSum.count(1500)); - - // Advance by 2 buckets (from 1000 to 3000) - assertEquals(15, movingSum.sum(3000)); - movingSum.add(3500, 20); - assertEquals(35, movingSum.sum(3500)); - assertEquals(3, movingSum.count(3500)); - - // Wait 11 seconds (moving completely outside window) - assertEquals(0, movingSum.sum(12000)); - assertEquals(0, movingSum.count(12000)); - assertFalse(movingSum.hasData(12000)); - } - - @Test - public void testInvalidArguments() { - assertThrows(IllegalArgumentException.class, () -> new MovingSum(100, 1000)); - assertThrows(IllegalArgumentException.class, () -> new MovingSum(1000, 0)); - } -} diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle index 2686297c0001..c8cc9958e2f3 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 15c4e678bcb4..000000000000 --- 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 c9507475648d..6966de5b484b 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.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 6db8af3906c3..000000000000 --- 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)); - } - } - } -} From 4b75f566371a7a2a8b88cea72986c498aabcf3df Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 1 Jul 2026 12:36:10 -0400 Subject: [PATCH 7/7] Remove empty util directory --- .../sdk/io/components/util/package-info.java | 20 ------------------- 1 file changed, 20 deletions(-) delete mode 100644 sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java deleted file mode 100644 index 40c757128789..000000000000 --- a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/util/package-info.java +++ /dev/null @@ -1,20 +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. - */ - -/** Defines utilities for Beam IO components. */ -package org.apache.beam.sdk.io.components.util;