This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new 2d7be922ceb4 fix(flink): prevent write rate limiter stalls and thread
leaks (#20189)
2d7be922ceb4 is described below
commit 2d7be922ceb4bb671847b5abc687dc30ba0b8112
Author: fhan <[email protected]>
AuthorDate: Tue Oct 13 10:56:11 2026 +0800
fix(flink): prevent write rate limiter stalls and thread leaks (#20189)
* fix(flink): prevent write rate limiter stalls and thread leaks
---------
Co-authored-by: fhan <[email protected]>
---
.../org/apache/hudi/common/util/RateLimiter.java | 108 +++++++++++++++++---
.../apache/hudi/common/util/TestRatelimiter.java | 110 ++++++++++++++++++---
.../append/AppendWriteFunctionWithRateLimit.java | 26 +++--
.../RowDataToHoodieFunctionWithRateLimit.java | 26 +++--
.../hudi/sink/utils/FlinkRateLimiterFactory.java | 92 +++++++++++++++++
.../TestAppendWriteFunctionWithRateLimit.java | 50 ++++++++++
.../TestRowDataToHoodieFunctionWithRateLimit.java | 68 +++++++++++++
.../sink/utils/TestFlinkRateLimiterFactory.java | 80 +++++++++++++++
8 files changed, 522 insertions(+), 38 deletions(-)
diff --git
a/hudi-common/src/main/java/org/apache/hudi/common/util/RateLimiter.java
b/hudi-common/src/main/java/org/apache/hudi/common/util/RateLimiter.java
index ab31d5013022..978b6fa21951 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/util/RateLimiter.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/util/RateLimiter.java
@@ -36,25 +36,67 @@ public class RateLimiter {
private final Semaphore semaphore;
private final int maxPermits;
+ private final long releasePermitsPeriod;
private final TimeUnit timePeriod;
+ private final Runnable permitWaitCallback;
+ private volatile boolean stopped;
private ScheduledExecutorService scheduler;
private static final long RELEASE_PERMITS_PERIOD_IN_SECONDS = 1L;
private static final long WAIT_BEFORE_NEXT_ACQUIRE_PERMIT_IN_MS = 5;
private static final int SCHEDULER_CORE_THREAD_POOL_SIZE = 1;
public static RateLimiter create(int permits, TimeUnit timePeriod) {
- final RateLimiter limiter = new RateLimiter(permits, timePeriod);
+ return create(permits, RELEASE_PERMITS_PERIOD_IN_SECONDS, timePeriod);
+ }
+
+ public static RateLimiter create(int permits, long releasePermitsPeriod,
TimeUnit timePeriod) {
+ validateConfiguration(permits, releasePermitsPeriod, timePeriod);
+ final RateLimiter limiter = new RateLimiter(permits, releasePermitsPeriod,
timePeriod);
+ limiter.releasePermitsPeriodically();
+ return limiter;
+ }
+
+ @VisibleForTesting
+ public static RateLimiter create(int permits, long releasePermitsPeriod,
TimeUnit timePeriod,
+ ScheduledExecutorService scheduler) {
+ validateConfiguration(permits, releasePermitsPeriod, timePeriod);
+ ValidationUtils.checkArgument(scheduler != null, "Scheduler must not be
null");
+ final RateLimiter limiter = new RateLimiter(permits, releasePermitsPeriod,
timePeriod);
+ limiter.releasePermitsPeriodically(scheduler);
+ return limiter;
+ }
+
+ @VisibleForTesting
+ static RateLimiter create(int permits, long releasePermitsPeriod, TimeUnit
timePeriod,
+ Runnable permitWaitCallback) {
+ validateConfiguration(permits, releasePermitsPeriod, timePeriod);
+ final RateLimiter limiter = new RateLimiter(
+ permits, releasePermitsPeriod, timePeriod, permitWaitCallback);
limiter.releasePermitsPeriodically();
return limiter;
}
- private RateLimiter(int permits, TimeUnit timePeriod) {
+ private static void validateConfiguration(int permits, long
releasePermitsPeriod, TimeUnit timePeriod) {
+ ValidationUtils.checkArgument(permits > 0, "Permits must be greater than
zero");
+ ValidationUtils.checkArgument(releasePermitsPeriod > 0, "Release permits
period must be greater than zero");
+ ValidationUtils.checkArgument(timePeriod != null, "Time period must not be
null");
+ }
+
+ private RateLimiter(int permits, long releasePermitsPeriod, TimeUnit
timePeriod) {
+ this(permits, releasePermitsPeriod, timePeriod, null);
+ }
+
+ private RateLimiter(int permits, long releasePermitsPeriod, TimeUnit
timePeriod,
+ Runnable permitWaitCallback) {
this.semaphore = new Semaphore(permits);
this.maxPermits = permits;
+ this.releasePermitsPeriod = releasePermitsPeriod;
this.timePeriod = timePeriod;
+ this.permitWaitCallback = permitWaitCallback;
}
public boolean tryAcquire(int numPermits) {
+ ValidationUtils.checkArgument(numPermits > 0, "Number of permits must be
greater than zero");
int remainingPermits = numPermits;
while (remainingPermits > 0) {
if (remainingPermits > maxPermits) {
@@ -68,27 +110,71 @@ public class RateLimiter {
}
public boolean acquire(int numOps) {
+ ValidationUtils.checkArgument(numOps > 0 && numOps <= maxPermits,
+ "Number of permits must be between one and the configured maximum");
+ return acquireInternal(numOps);
+ }
+
+ public boolean acquire() {
+ return acquireInternal(1);
+ }
+
+ private boolean acquireInternal(int numOps) {
try {
- while (!semaphore.tryAcquire(numOps)) {
+ while (!stopped && !semaphore.tryAcquire(numOps)) {
+ if (permitWaitCallback != null) {
+ permitWaitCallback.run();
+ }
Thread.sleep(WAIT_BEFORE_NEXT_ACQUIRE_PERMIT_IN_MS);
}
+ ValidationUtils.checkState(!stopped, "Rate limiter is stopped");
log.debug("acquire permits: {}, maxPermits: {}", numOps, maxPermits);
} catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
throw new RuntimeException("Unable to acquire permits", e);
}
return true;
}
- public void stop() {
- scheduler.shutdownNow();
+ public synchronized void stop() {
+ if (stopped) {
+ return;
+ }
+ stopped = true;
+ if (scheduler != null) {
+ scheduler.shutdownNow();
+ }
+ }
+
+ public boolean isStopped() {
+ return stopped;
}
- public void releasePermitsPeriodically() {
- scheduler =
Executors.newScheduledThreadPool(SCHEDULER_CORE_THREAD_POOL_SIZE);
- scheduler.scheduleAtFixedRate(() -> {
- log.debug("Release permits: maxPermits: {}, available: {}", maxPermits,
semaphore.availablePermits());
- semaphore.release(maxPermits - semaphore.availablePermits());
- }, RELEASE_PERMITS_PERIOD_IN_SECONDS, RELEASE_PERMITS_PERIOD_IN_SECONDS,
timePeriod);
+ public synchronized void releasePermitsPeriodically() {
+ ValidationUtils.checkState(!stopped, "Cannot start a stopped rate
limiter");
+ if (scheduler != null) {
+ return;
+ }
+
releasePermitsPeriodically(Executors.newScheduledThreadPool(SCHEDULER_CORE_THREAD_POOL_SIZE,
+ new CustomizedThreadFactory("rate-limiter", true)));
+ }
+
+ private synchronized void
releasePermitsPeriodically(ScheduledExecutorService scheduler) {
+ ValidationUtils.checkState(!stopped, "Cannot start a stopped rate
limiter");
+ if (this.scheduler != null) {
+ return;
+ }
+ this.scheduler = scheduler;
+ try {
+ scheduler.scheduleAtFixedRate(() -> {
+ log.debug("Release permits: maxPermits: {}, available: {}",
maxPermits, semaphore.availablePermits());
+ semaphore.release(maxPermits - semaphore.availablePermits());
+ }, releasePermitsPeriod, releasePermitsPeriod, timePeriod);
+ } catch (RuntimeException e) {
+ this.scheduler = null;
+ scheduler.shutdownNow();
+ throw e;
+ }
}
diff --git
a/hudi-common/src/test/java/org/apache/hudi/common/util/TestRatelimiter.java
b/hudi-common/src/test/java/org/apache/hudi/common/util/TestRatelimiter.java
index 3c27138bf2d5..d08f7f4fc039 100644
--- a/hudi-common/src/test/java/org/apache/hudi/common/util/TestRatelimiter.java
+++ b/hudi-common/src/test/java/org/apache/hudi/common/util/TestRatelimiter.java
@@ -19,11 +19,23 @@
package org.apache.hudi.common.util;
import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
/**
* Tests {@link RateLimiter}.
@@ -33,25 +45,93 @@ public class TestRatelimiter {
@Test
public void testRateLimiterWithNoThrottling() throws InterruptedException {
RateLimiter limiter = RateLimiter.create(1000, TimeUnit.SECONDS);
- long start = System.currentTimeMillis();
- assertEquals(true, limiter.tryAcquire(1000));
- // Sleep to represent some operation
- Thread.sleep(500);
- long end = System.currentTimeMillis();
- // With a large permit limit, there shouldn't be any throttling of
operations
- assertTrue((end - start) < TimeUnit.SECONDS.toMillis(2));
+ try {
+ long start = System.currentTimeMillis();
+ assertEquals(true, limiter.tryAcquire(1000));
+ // Sleep to represent some operation
+ Thread.sleep(500);
+ long end = System.currentTimeMillis();
+ // With a large permit limit, there shouldn't be any throttling of
operations
+ assertTrue((end - start) < TimeUnit.SECONDS.toMillis(2));
+ } finally {
+ limiter.stop();
+ }
}
@Test
public void testRateLimiterWithThrottling() throws InterruptedException {
RateLimiter limiter = RateLimiter.create(100, TimeUnit.SECONDS);
- long start = System.currentTimeMillis();
- assertEquals(true, limiter.tryAcquire(400));
- // Sleep to represent some operation
- Thread.sleep(500);
- long end = System.currentTimeMillis();
- // As size of operations is more than the maximum permits per second,
- // whole execution should be greater than 1 second
- assertTrue((end - start) >= TimeUnit.SECONDS.toMillis(2));
+ try {
+ long start = System.currentTimeMillis();
+ assertEquals(true, limiter.tryAcquire(400));
+ // Sleep to represent some operation
+ Thread.sleep(500);
+ long end = System.currentTimeMillis();
+ // As size of operations is more than the maximum permits per second,
+ // whole execution should be greater than 1 second
+ assertTrue((end - start) >= TimeUnit.SECONDS.toMillis(2));
+ } finally {
+ limiter.stop();
+ }
+ }
+
+ @Test
+ public void testCustomReleasePeriod() {
+ ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
+
+ RateLimiter limiter = RateLimiter.create(1, 100, TimeUnit.MILLISECONDS,
scheduler);
+ ArgumentCaptor<Runnable> refillTask =
ArgumentCaptor.forClass(Runnable.class);
+ try {
+ verify(scheduler).scheduleAtFixedRate(
+ refillTask.capture(), eq(100L), eq(100L), eq(TimeUnit.MILLISECONDS));
+ assertTrue(limiter.acquire());
+ refillTask.getValue().run();
+ assertTrue(limiter.acquire());
+ } finally {
+ limiter.stop();
+ }
+ verify(scheduler).shutdownNow();
+ }
+
+ @Test
+ public void testInvalidConfigurationIsRejected() {
+ assertThrows(IllegalArgumentException.class, () -> RateLimiter.create(0,
TimeUnit.SECONDS));
+ assertThrows(IllegalArgumentException.class, () -> RateLimiter.create(1,
0, TimeUnit.SECONDS));
+ assertThrows(IllegalArgumentException.class, () -> RateLimiter.create(1,
1, null));
+ }
+
+ @Test
+ public void testInvalidAcquireRequestIsRejected() {
+ RateLimiter limiter = RateLimiter.create(1, 1, TimeUnit.DAYS);
+ try {
+ assertThrows(IllegalArgumentException.class, () -> limiter.acquire(0));
+ assertThrows(IllegalArgumentException.class, () -> limiter.acquire(2));
+ } finally {
+ limiter.stop();
+ }
+ }
+
+ @Test
+ public void testStopIsIdempotentAndUnblocksAcquire() throws Exception {
+ CountDownLatch waitingForPermit = new CountDownLatch(1);
+ RateLimiter limiter = RateLimiter.create(1, 1, TimeUnit.DAYS,
waitingForPermit::countDown);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ assertTrue(limiter.acquire());
+ Future<Boolean> blockedAcquire = executor.submit(() ->
limiter.acquire());
+ assertTrue(waitingForPermit.await(1, TimeUnit.SECONDS));
+
+ limiter.stop();
+ limiter.stop();
+
+ ExecutionException exception = assertThrows(
+ ExecutionException.class, () -> blockedAcquire.get(1,
TimeUnit.SECONDS));
+ assertInstanceOf(IllegalStateException.class, exception.getCause());
+ assertTrue(limiter.isStopped());
+ } finally {
+ limiter.stop();
+ executor.shutdownNow();
+ assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS));
+ }
}
}
diff --git
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithRateLimit.java
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithRateLimit.java
index 00d70a7b417f..977829d46937 100644
---
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithRateLimit.java
+++
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/append/AppendWriteFunctionWithRateLimit.java
@@ -20,6 +20,7 @@ package org.apache.hudi.sink.append;
import org.apache.hudi.common.util.RateLimiter;
import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.sink.utils.FlinkRateLimiterFactory;
import org.apache.hudi.utils.RuntimeContextUtils;
import org.apache.flink.configuration.Configuration;
@@ -27,8 +28,6 @@ import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.Collector;
-import java.util.concurrent.TimeUnit;
-
/**
* Append write function with configurable rate limit.
*/
@@ -37,7 +36,7 @@ public class AppendWriteFunctionWithRateLimit<I>
/**
* Total rate limit per second for this job.
*/
- private final double totalLimit;
+ private final long totalLimit;
/**
* Rate limit per second for per task.
@@ -52,13 +51,28 @@ public class AppendWriteFunctionWithRateLimit<I>
@Override
public void open(Configuration parameters) throws Exception {
super.open(parameters);
- this.rateLimiter =
- RateLimiter.create((int) totalLimit /
RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()),
TimeUnit.SECONDS);
+ this.rateLimiter = FlinkRateLimiterFactory.create(
+ totalLimit,
RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()));
}
@Override
public void processElement(I value, Context ctx, Collector<RowData> out)
throws Exception {
- rateLimiter.acquire(1);
+ rateLimiter.acquire();
super.processElement(value, ctx, out);
}
+
+ @Override
+ public void close() throws Exception {
+ try {
+ super.close();
+ } finally {
+ if (rateLimiter != null) {
+ rateLimiter.stop();
+ }
+ }
+ }
+
+ RateLimiter getRateLimiter() {
+ return rateLimiter;
+ }
}
diff --git
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/transform/RowDataToHoodieFunctionWithRateLimit.java
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/transform/RowDataToHoodieFunctionWithRateLimit.java
index 2cd01da03a2d..673817c8c83e 100644
---
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/transform/RowDataToHoodieFunctionWithRateLimit.java
+++
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/transform/RowDataToHoodieFunctionWithRateLimit.java
@@ -21,14 +21,13 @@ package org.apache.hudi.sink.transform;
import org.apache.hudi.client.model.HoodieFlinkInternalRow;
import org.apache.hudi.common.util.RateLimiter;
import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.sink.utils.FlinkRateLimiterFactory;
import org.apache.hudi.utils.RuntimeContextUtils;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;
-import java.util.concurrent.TimeUnit;
-
/**
* Function that transforms RowData to a {@code HoodieFlinkInternalRow} with
RateLimit.
*/
@@ -37,7 +36,7 @@ public class RowDataToHoodieFunctionWithRateLimit<I extends
RowData, O extends H
/**
* Total rate limit per second for this job.
*/
- private final double totalLimit;
+ private final long totalLimit;
/**
* Rate limit per second for per task.
@@ -52,13 +51,28 @@ public class RowDataToHoodieFunctionWithRateLimit<I extends
RowData, O extends H
@Override
public void open(Configuration parameters) throws Exception {
super.open(parameters);
- this.rateLimiter =
- RateLimiter.create((int) totalLimit /
RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()),
TimeUnit.SECONDS);
+ this.rateLimiter = FlinkRateLimiterFactory.create(
+ totalLimit,
RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()));
}
@Override
public O map(I i) throws Exception {
- rateLimiter.acquire(1);
+ rateLimiter.acquire();
return super.map(i);
}
+
+ @Override
+ public void close() throws Exception {
+ try {
+ super.close();
+ } finally {
+ if (rateLimiter != null) {
+ rateLimiter.stop();
+ }
+ }
+ }
+
+ RateLimiter getRateLimiter() {
+ return rateLimiter;
+ }
}
diff --git
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/FlinkRateLimiterFactory.java
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/FlinkRateLimiterFactory.java
new file mode 100644
index 000000000000..54ca23dbb8b4
--- /dev/null
+++
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/FlinkRateLimiterFactory.java
@@ -0,0 +1,92 @@
+/*
+ * 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.hudi.sink.utils;
+
+import org.apache.hudi.common.util.RateLimiter;
+import org.apache.hudi.common.util.ValidationUtils;
+import org.apache.hudi.common.util.VisibleForTesting;
+
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Factory for distributing a job-wide write rate limit across Flink subtasks.
+ */
+public final class FlinkRateLimiterFactory {
+
+ private FlinkRateLimiterFactory() {
+ }
+
+ public static RateLimiter create(long totalRateLimit, int parallelism) {
+ RateLimitConfig rateLimitConfig = resolveRateLimit(totalRateLimit,
parallelism);
+ return RateLimiter.create(
+ rateLimitConfig.getPermits(), rateLimitConfig.getReleasePeriod(),
rateLimitConfig.getTimeUnit());
+ }
+
+ @VisibleForTesting
+ static RateLimiter create(long totalRateLimit, int parallelism,
ScheduledExecutorService scheduler) {
+ RateLimitConfig rateLimitConfig = resolveRateLimit(totalRateLimit,
parallelism);
+ return RateLimiter.create(
+ rateLimitConfig.getPermits(), rateLimitConfig.getReleasePeriod(),
rateLimitConfig.getTimeUnit(), scheduler);
+ }
+
+ static RateLimitConfig resolveRateLimit(long totalRateLimit, int
parallelism) {
+ ValidationUtils.checkArgument(totalRateLimit > 0, "Total rate limit must
be greater than zero");
+ ValidationUtils.checkArgument(parallelism > 0, "Parallelism must be
greater than zero");
+
+ if (totalRateLimit >= parallelism) {
+ long rateLimitPerSubtask = totalRateLimit / parallelism;
+ ValidationUtils.checkArgument(rateLimitPerSubtask <= Integer.MAX_VALUE,
+ "Rate limit per subtask exceeds the supported maximum");
+ return new RateLimitConfig((int) rateLimitPerSubtask, 1,
TimeUnit.SECONDS);
+ }
+
+ long periodNanos = divideRoundingUp(
+ Math.multiplyExact(TimeUnit.SECONDS.toNanos(1), parallelism),
totalRateLimit);
+ return new RateLimitConfig(1, periodNanos, TimeUnit.NANOSECONDS);
+ }
+
+ private static long divideRoundingUp(long dividend, long divisor) {
+ return 1 + (dividend - 1) / divisor;
+ }
+
+ static final class RateLimitConfig {
+ private final int permits;
+ private final long releasePeriod;
+ private final TimeUnit timeUnit;
+
+ private RateLimitConfig(int permits, long releasePeriod, TimeUnit
timeUnit) {
+ this.permits = permits;
+ this.releasePeriod = releasePeriod;
+ this.timeUnit = timeUnit;
+ }
+
+ int getPermits() {
+ return permits;
+ }
+
+ long getReleasePeriod() {
+ return releasePeriod;
+ }
+
+ TimeUnit getTimeUnit() {
+ return timeUnit;
+ }
+ }
+}
diff --git
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithRateLimit.java
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithRateLimit.java
new file mode 100644
index 000000000000..980e65a1ac76
--- /dev/null
+++
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/append/TestAppendWriteFunctionWithRateLimit.java
@@ -0,0 +1,50 @@
+/*
+ * 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.hudi.sink.append;
+
+import org.apache.hudi.common.util.RateLimiter;
+import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.sink.utils.MockStreamingRuntimeContext;
+import org.apache.hudi.utils.TestConfigurations;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.data.RowData;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TestAppendWriteFunctionWithRateLimit {
+
+ @Test
+ void testCloseStopsRateLimiter() throws Exception {
+ Configuration config = new Configuration();
+ config.set(FlinkOptions.WRITE_RATE_LIMIT, 1L);
+ AppendWriteFunctionWithRateLimit<RowData> function =
+ new AppendWriteFunctionWithRateLimit<>(TestConfigurations.ROW_TYPE,
config);
+ function.setRuntimeContext(new MockStreamingRuntimeContext(false, 4, 0));
+ function.open(new Configuration());
+ RateLimiter limiter = function.getRateLimiter();
+
+ assertFalse(limiter.isStopped());
+ function.close();
+ function.close();
+ assertTrue(limiter.isStopped());
+ }
+}
diff --git
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/transform/TestRowDataToHoodieFunctionWithRateLimit.java
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/transform/TestRowDataToHoodieFunctionWithRateLimit.java
new file mode 100644
index 000000000000..6d86b0f28da8
--- /dev/null
+++
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/transform/TestRowDataToHoodieFunctionWithRateLimit.java
@@ -0,0 +1,68 @@
+/*
+ * 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.hudi.sink.transform;
+
+import org.apache.hudi.client.model.HoodieFlinkInternalRow;
+import org.apache.hudi.common.util.RateLimiter;
+import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.sink.utils.MockStreamingRuntimeContext;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.table.types.logical.RowType;
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TestRowDataToHoodieFunctionWithRateLimit {
+
+ @Test
+ void testLowRateLimitDoesNotBlockAndCloseStopsLimiter() throws Exception {
+ Configuration config = new Configuration();
+ config.set(FlinkOptions.WRITE_RATE_LIMIT, 1L);
+ config.set(FlinkOptions.RECORD_KEY_FIELD, "id");
+ config.set(FlinkOptions.PARTITION_PATH_FIELD, "partition");
+ RowType rowType = (RowType) DataTypes.ROW(
+ DataTypes.FIELD("id", DataTypes.STRING()),
+ DataTypes.FIELD("partition", DataTypes.STRING())).getLogicalType();
+ RowDataToHoodieFunctionWithRateLimit<RowData, HoodieFlinkInternalRow>
function =
+ new RowDataToHoodieFunctionWithRateLimit<>(rowType, config);
+ function.setRuntimeContext(new MockStreamingRuntimeContext(false, 4, 0));
+ function.open(new Configuration());
+ RateLimiter limiter = function.getRateLimiter();
+
+ try {
+ assertFalse(limiter.isStopped());
+ GenericRowData row = GenericRowData.of(
+ StringData.fromString("id-1"), StringData.fromString("p1"));
+ assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
function.map(row));
+ } finally {
+ function.close();
+ }
+
+ assertTrue(limiter.isStopped());
+ }
+}
diff --git
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFlinkRateLimiterFactory.java
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFlinkRateLimiterFactory.java
new file mode 100644
index 000000000000..75d4573674e9
--- /dev/null
+++
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/TestFlinkRateLimiterFactory.java
@@ -0,0 +1,80 @@
+/*
+ * 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.hudi.sink.utils;
+
+import org.apache.hudi.common.util.RateLimiter;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+class TestFlinkRateLimiterFactory {
+
+ @Test
+ void testRateLimitLowerThanParallelismUsesLongerReleasePeriod() {
+ ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
+ long expectedPeriodNanos = TimeUnit.SECONDS.toNanos(4);
+
+ RateLimiter limiter = FlinkRateLimiterFactory.create(1, 4, scheduler);
+ ArgumentCaptor<Runnable> refillTask =
ArgumentCaptor.forClass(Runnable.class);
+ try {
+ verify(scheduler).scheduleAtFixedRate(
+ refillTask.capture(), eq(expectedPeriodNanos),
eq(expectedPeriodNanos), eq(TimeUnit.NANOSECONDS));
+ assertTrue(limiter.acquire());
+ refillTask.getValue().run();
+ assertTrue(limiter.acquire());
+ } finally {
+ limiter.stop();
+ }
+ verify(scheduler).shutdownNow();
+ }
+
+ @Test
+ void testRateLimitAtLeastParallelismUsesPermitsPerSecond() {
+ FlinkRateLimiterFactory.RateLimitConfig config =
FlinkRateLimiterFactory.resolveRateLimit(8, 4);
+
+ assertEquals(2, config.getPermits());
+ assertEquals(1, config.getReleasePeriod());
+ assertEquals(TimeUnit.SECONDS, config.getTimeUnit());
+ }
+
+ @Test
+ void testFractionalRateLimitRoundsReleasePeriodUp() {
+ FlinkRateLimiterFactory.RateLimitConfig config =
FlinkRateLimiterFactory.resolveRateLimit(3, 4);
+
+ assertEquals(1, config.getPermits());
+ assertEquals(1_333_333_334L, config.getReleasePeriod());
+ assertEquals(TimeUnit.NANOSECONDS, config.getTimeUnit());
+ }
+
+ @Test
+ void testInvalidRateLimitInputsAreRejected() {
+ assertThrows(IllegalArgumentException.class, () ->
FlinkRateLimiterFactory.create(0, 4));
+ assertThrows(IllegalArgumentException.class, () ->
FlinkRateLimiterFactory.create(1, 0));
+ }
+}
