{@link RateLimiter#create(double)} builds a {@code SmoothBursty} limiter whose burst capacity + * is at most one second's worth of permits, and {@link RateLimiter#tryAcquire(int, long, TimeUnit)} + * only grants a request when its whole cost fits inside the given timeout. When a single request + * costs more than the timeout, every attempt fails instantly and the surrounding retry loop + * degenerates into a 100% CPU busy-wait that never logs anything. Such a request must therefore + * never be passed to {@code tryAcquire} in one piece, which is what this class takes care of. + */ +public class RateLimitUtils { + + private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitUtils.class); + + /** The default value of {@code RateLimiterHotReloadCheckIntervalMs}. */ + public static final long DEFAULT_CHECK_INTERVAL_MS = 1000; + + // The loop below is a busy-wait guard as well: a healthy chunk needs at most one timed attempt + // before it falls back to a blocking acquisition, so exceeding this budget means something is + // wrong (e.g. the limiter is stuck) rather than that we merely have to wait longer. + private static final int MAX_ATTEMPTS_PER_CHUNK = 20; + + private RateLimitUtils() { + // Utility class + } + + /** Decides whether to keep waiting for the rate limiter. */ + @FunctionalInterface + public interface ContinueChecker { + /** + * @return {@code true} to keep acquiring, {@code false} to abort the whole acquisition + */ + boolean shouldContinue(); + } + + /** + * Acquires {@code bytes} permits from {@code rateLimiter}, checking the limiter every {@link + * #DEFAULT_CHECK_INTERVAL_MS}. + * + * @param rateLimiter the permit-based limiter to acquire from + * @param bytes the number of bytes to acquire, non-positive values return immediately + * @param continueChecker consulted between pieces to stop waiting for a dead pipe + * @return {@code true} if the whole request was acquired, {@code false} if {@code + * continueChecker} asked to abort + */ + public static boolean acquireWithRateLimit( + final RateLimiter rateLimiter, final long bytes, final ContinueChecker continueChecker) { + return acquireWithRateLimit(rateLimiter, bytes, continueChecker, DEFAULT_CHECK_INTERVAL_MS); + } + + /** + * Acquires {@code bytes} permits from {@code rateLimiter}. + * + *
The request is split into pieces that the burst capacity can actually grant, each piece is + * acquired with a timeout that covers its own cost, and a piece that keeps timing out (which can + * only happen when the limiter's rate is reloaded downwards) falls back to a blocking acquisition + * so that the caller stops burning CPU. The average throughput always converges to the configured + * rate. + * + * @param rateLimiter the permit-based limiter to acquire from + * @param bytes the number of bytes to acquire, non-positive values return immediately + * @param continueChecker consulted between pieces and on a timeout, to stop waiting for a dead + * pipe + * @param checkIntervalMs how often {@code continueChecker} is consulted while waiting + * @return {@code true} if the whole request was acquired, {@code false} if {@code + * continueChecker} asked to abort + */ + public static boolean acquireWithRateLimit( + final RateLimiter rateLimiter, + final long bytes, + final ContinueChecker continueChecker, + final long checkIntervalMs) { + final long boundedCheckIntervalMs = Math.max(1L, checkIntervalMs); + + long remaining = bytes; + boolean logPipeSizeMismatch = true; + + while (remaining > 0) { + // A request that fits into a single piece is served by the limiter's initial burst without + // waiting, so the checker only matters for requests that need several pieces, i.e. the ones + // that actually take time to transfer. + if (remaining < bytes && !continueChecker.shouldContinue()) { + return false; + } + + final long piece = calculatePieceSize(rateLimiter, remaining); + if (logPipeSizeMismatch && piece < remaining) { + logPipeSizeMismatch = false; + LOGGER.warn( + PipeMessages + .LOG_PIPE_RATE_LIMIT_IS_SMALLER_THAN_THE_SINGLE_REQUEST_SIZE_SO_ARG_BYTES_WILL_BE_TRANSFERRED_IN_ARG_SIZED_PIECES_AT_ARG_BYTES_PER_SECOND_CONSIDER_SETTING_A_LARGER_RATE_LIMIT_FOR_BETTER_PERFORMANCE_4450D4DC, + bytes, + piece, + rateLimiter.getRate()); + } + + final long timeoutMs = calculateTimeoutMs(rateLimiter, piece, boundedCheckIntervalMs); + int attempts = 0; + boolean acquired = false; + while (!acquired) { + acquired = rateLimiter.tryAcquire((int) piece, timeoutMs, TimeUnit.MILLISECONDS); + if (acquired) { + break; + } + + if (!continueChecker.shouldContinue()) { + return false; + } + + if (++attempts >= MAX_ATTEMPTS_PER_CHUNK) { + // The rate has most likely been reloaded downwards while we were waiting, so our + // estimate is far too optimistic. Wait for the permits without a deadline instead of + // retrying with a deadline that we now know we cannot meet, which would only spin. + rateLimiter.acquire((int) piece); + acquired = true; + } + } + + remaining -= piece; + } + + return true; + } + + /** + * @return a piece size that is small enough for the limiter's burst capacity to grant + */ + private static long calculatePieceSize(final RateLimiter rateLimiter, final long bytes) { + // The burst capacity of RateLimiter.create(rate) is one second's worth of permits, so anything + // larger than that can never be granted by a single timeout-bounded attempt. + final double bytesPerSecond = rateLimiter.getRate(); + if (!(bytesPerSecond > 0) || bytesPerSecond >= Integer.MAX_VALUE) { + // Rate limiting is effectively disabled (e.g. the rate is set to Double.MAX_VALUE), but the + // piece still has to stay within the int range of the limiter's API. + return Math.min(bytes, Integer.MAX_VALUE); + } + return Math.min(bytes, (long) bytesPerSecond); + } + + /** + * @return a timeout that covers the expected wait for {@code piece}, so that a legitimate wait is + * not mistaken for a timeout + */ + private static long calculateTimeoutMs( + final RateLimiter rateLimiter, final long piece, final long checkIntervalMs) { + final double rate = rateLimiter.getRate(); + if (!(rate > 0)) { + return checkIntervalMs; + } + + final double expectedWaitMs = piece / rate * 1000.0; + if (!(expectedWaitMs > 0) || expectedWaitMs >= Long.MAX_VALUE - checkIntervalMs) { + return checkIntervalMs; + } + + return Math.max(checkIntervalMs, (long) expectedWaitMs + checkIntervalMs); + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/PipeRateLimiterFieldScenarioTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/PipeRateLimiterFieldScenarioTest.java new file mode 100644 index 0000000000000..f20bae180dc41 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/PipeRateLimiterFieldScenarioTest.java @@ -0,0 +1,136 @@ +/* + * 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.iotdb.commons.pipe.sink.limiter; + +import org.apache.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.pipe.config.PipeConfig; + +import com.google.common.util.concurrent.RateLimiter; +import org.junit.Assert; +import org.junit.Test; + +import java.lang.management.ManagementFactory; +import java.lang.management.ThreadMXBean; +import java.util.concurrent.TimeUnit; + +/** + * Reproduces the configuration of the pipe regression case {@code + * obj-10-table-big-object-kill-restart}: a {@code rate-limit-bytes-per-second} of 1 MB/s together + * with an OBJECT batch that is as large as {@code PipeSinkReadFileBufferSize} (5 MB by default). + * + *
With those two defaults the cost of one batch (5 seconds at 1 MB/s) is larger than the + * hot-reload check interval that used to double as the {@code tryAcquire} timeout (1 second), so + * every attempt failed instantly and the sink thread spun at roughly 100% CPU while the pipe kept + * reporting {@code RUNNING} with an empty {@code ExceptionMessage}. The stall was silent because + * that retry loop logged nothing. + * + *
Such a request is never rejected outright - the quota keeps accumulating during the spin and + * the transfer does make progress at the configured rate - so the defect shows up as wasted CPU + * rather than as a wrong result. These tests therefore assert on the CPU time consumed while the + * requests are being paced. + */ +public class PipeRateLimiterFieldScenarioTest { + + /** The {@code sink('rate-limit-bytes-per-second')} used by the regression case. */ + private static final double RATE_LIMIT_BYTES_PER_SECOND = 1048576; + + /** One OBJECT batch is bounded by {@code PipeSinkReadFileBufferSize}, 5 MB by default. */ + private static final int OBJECT_BATCH_BYTES = 5 * 1024 * 1024; + + /** + * A spin consumes CPU for almost the whole waiting time, whereas a correct implementation sleeps + * and consumes almost none. Measured with the pre-fix logic the ratio is about 64%, and with the + * fix it is 0%, so this threshold separates the two with a wide margin. + */ + private static final double MAX_CPU_TIME_RATIO = 0.2; + + private static final ThreadMXBean THREAD_MX_BEAN = ManagementFactory.getThreadMXBean(); + + @Test(timeout = 120_000) + public void testFieldScenarioDoesNotBusyWait() { + final RateLimiter rateLimiter = RateLimiter.create(RATE_LIMIT_BYTES_PER_SECOND); + + // Three batches, i.e. 15 MB, so that the initial one-second burst no longer deserves the + // transfer and the limiter has to pace the remaining ones. The complete object is 1 GB, which + // would take about 17 minutes at this rate and cannot be part of a unit test. + final long wallTimeNanos = System.nanoTime(); + final long cpuTimeNanos = THREAD_MX_BEAN.getCurrentThreadCpuTime(); + for (int batch = 0; batch < 3; batch++) { + Assert.assertTrue( + RateLimitUtils.acquireWithRateLimit( + rateLimiter, + OBJECT_BATCH_BYTES, + () -> true, + PipeConfig.getInstance().getRateLimiterHotReloadCheckIntervalMs())); + } + final long elapsedNanos = System.nanoTime() - wallTimeNanos; + final long consumedCpuNanos = THREAD_MX_BEAN.getCurrentThreadCpuTime() - cpuTimeNanos; + + final double elapsedSeconds = elapsedNanos / 1_000_000_000.0; + final double consumedCpuSeconds = consumedCpuNanos / 1_000_000_000.0; + final double cpuTimeRatio = consumedCpuSeconds / elapsedSeconds; + + // The transfer must be paced by the rate limit instead of being served by the burst, which is + // the + // case exactly when the limiter is still being asked for one batch at a time. + Assert.assertTrue( + "15 MB at 1 MB/s minus the initial burst should take at least 12s, but took " + + elapsedSeconds + + "s", + elapsedSeconds >= 12.0); + Assert.assertTrue( + "the CPU time ratio is " + + cpuTimeRatio + + " (" + + consumedCpuSeconds + + "s CPU over " + + elapsedSeconds + + "s wall), which means the sink thread was busy-waiting instead of waiting for the" + + " rate limiter", + cpuTimeRatio <= MAX_CPU_TIME_RATIO); + } + + /** + * The burst capacity of {@code RateLimiter.create(rate)} is one second's worth of permits, so a + * single batch must stay below it for one timeout-bounded attempt to be able to grant the whole + * batch. + */ + @Test + public void testConfiguredDefaultsStayWithinTheBurstCapacity() { + final CommonConfig config = CommonDescriptor.getInstance().getConfig(); + + final long burstCapacity = (long) RATE_LIMIT_BYTES_PER_SECOND; + final long timeoutNanos = + TimeUnit.MILLISECONDS.toNanos(config.getRateLimiterHotReloadCheckIntervalMs()); + final long grantableBytes = burstCapacity * timeoutNanos / 1_000_000_000L; + + // A batch larger than this cannot be granted by a timeout-bounded attempt, which is the trap + // the + // regression case fell into. + Assert.assertTrue( + "one OBJECT batch (" + + OBJECT_BATCH_BYTES + + " bytes) exceeds what the rate limiter can grant within the check interval (" + + grantableBytes + + " bytes), so the limiter has to split the request", + OBJECT_BATCH_BYTES > grantableBytes); + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/RateLimitUtilsTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/RateLimitUtilsTest.java new file mode 100644 index 0000000000000..211fb90db48e7 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/limiter/RateLimitUtilsTest.java @@ -0,0 +1,150 @@ +/* + * 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.iotdb.commons.pipe.sink.limiter; + +import com.google.common.util.concurrent.RateLimiter; +import org.junit.Assert; +import org.junit.Test; + +import java.util.concurrent.atomic.AtomicBoolean; + +public class RateLimitUtilsTest { + + /** The rate limit and the request size reported from the field reproduction. */ + private static final double ONE_MB_PER_SECOND = 1024 * 1024; + + private static final long FIVE_MB = 5 * 1024 * 1024; + + private static final long CHECK_INTERVAL_MS = 1000; + + private static boolean acquire(final RateLimiter rateLimiter, final long bytes) { + return RateLimitUtils.acquireWithRateLimit(rateLimiter, bytes, () -> true, CHECK_INTERVAL_MS); + } + + /** + * A single request larger than the limiter's burst capacity (which is one second's worth of + * permits) must still be acquired. Before the fix, such a request made every attempt time out + * immediately, so the caller spun at 100% CPU without ever making progress. + */ + @Test(timeout = 60_000) + public void testAcquireRequestLargerThanBurstCapacity() { + final RateLimiter rateLimiter = RateLimiter.create(ONE_MB_PER_SECOND); + + final long startTime = System.nanoTime(); + final boolean acquired = acquire(rateLimiter, FIVE_MB); + final double elapsedSeconds = (System.nanoTime() - startTime) / 1_000_000_000.0; + + Assert.assertTrue(acquired); + // 5 MB at 1 MB/s, minus the one-second initial burst, must take about 4 seconds. + Assert.assertTrue( + "expected to be paced by the rate limiter, but took " + elapsedSeconds + "s", + elapsedSeconds >= 3.0); + } + + /** The average throughput must converge to the configured rate limit. */ + @Test(timeout = 60_000) + public void testAcquireIsPacedAtTheConfiguredRate() { + final RateLimiter rateLimiter = RateLimiter.create(FIVE_MB / 10.0); + + final long startTime = System.nanoTime(); + acquire(rateLimiter, FIVE_MB); + final double elapsedSeconds = (System.nanoTime() - startTime) / 1_000_000_000.0; + + // 5 MB at 0.5 MB/s is about 10 seconds, minus the one-second initial burst. + Assert.assertTrue( + "expected to be paced at the configured rate, but took " + elapsedSeconds + "s", + elapsedSeconds >= 7.0); + } + + /** A request of more than Integer.MAX_VALUE bytes must not overflow nor spin. */ + @Test(timeout = 120_000) + public void testAcquireRequestLargerThanIntegerMaxValue() { + final long bytes = 3L * Integer.MAX_VALUE; + + final RateLimiter rateLimiter = RateLimiter.create(256 * ONE_MB_PER_SECOND); + final long startTime = System.nanoTime(); + + Assert.assertTrue(acquire(rateLimiter, bytes)); + Assert.assertTrue( + "3 * Integer.MAX_VALUE bytes should be acquired in a few seconds", + (System.nanoTime() - startTime) / 1_000_000_000.0 < 60.0); + } + + /** + * The same must hold when rate limiting is disabled via {@code Double.MAX_VALUE}, even though + * such a rate makes the piece size calculation take its dedicated branch. + */ + @Test(timeout = 60_000) + public void testAcquireRequestLargerThanIntegerMaxValueWithRateLimitingDisabled() { + final RateLimiter rateLimiter = RateLimiter.create(Double.MAX_VALUE); + final long startTime = System.nanoTime(); + + // More than two pieces are needed, and every piece but the first is charged to the unlimited + // burst capacity, so this must complete without waiting. + Assert.assertTrue(acquire(rateLimiter, 3L * Integer.MAX_VALUE)); + Assert.assertTrue( + "an unlimited limiter should not make the caller wait", + (System.nanoTime() - startTime) / 1_000_000_000.0 < 10.0); + } + + /** + * When the checker asks to stop (e.g. the pipe has been dropped), a multi-piece request must + * abort instead of waiting for the remaining pieces. + */ + @Test(timeout = 60_000) + public void testAcquireAbortsWhenCheckerSaysStop() { + // The first piece is served by the initial burst, so the checker is consulted when the second + // one is about to be acquired and the request is dropped there. + final RateLimiter rateLimiter = RateLimiter.create(ONE_MB_PER_SECOND); + final AtomicBoolean checked = new AtomicBoolean(false); + + final long startTime = System.nanoTime(); + final boolean acquired = + RateLimitUtils.acquireWithRateLimit( + rateLimiter, + 4 * 1024 * 1024, + () -> { + checked.set(true); + return false; + }, + CHECK_INTERVAL_MS); + final double elapsedSeconds = (System.nanoTime() - startTime) / 1_000_000_000.0; + + Assert.assertTrue(checked.get()); + Assert.assertFalse(acquired); + // 4 MB at 1 MB/s would need about 3 seconds if the whole request were awaited. + Assert.assertTrue( + "expected to abort before waiting for the remaining pieces, but took " + + elapsedSeconds + + "s", + elapsedSeconds < 1.0); + } + + /** Non-positive requests and unlimited limiters are passed through. */ + @Test + public void testAcquireEdgeCases() { + final RateLimiter rateLimiter = RateLimiter.create(ONE_MB_PER_SECOND); + + Assert.assertTrue(acquire(rateLimiter, 0)); + Assert.assertTrue(acquire(rateLimiter, -1)); + // Double.MAX_VALUE is used to disable rate limiting, so nothing needs to be chunked or waited. + Assert.assertTrue(acquire(RateLimiter.create(Double.MAX_VALUE), FIVE_MB)); + } +}