This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-12099-def3d4bc635aaedc80123552adf99431bb5aac75 in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit 934276b2f735349c166569de58b490e21cc8319b Author: David Zollo <[email protected]> AuthorDate: Mon Oct 19 12:26:49 2026 +0000 [Fix][Connector-V2][RocketMQ] Close consumer client before interrupting its executor (#12099) Co-authored-by: davidzollo <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]> Co-authored-by: David Zollo <[email protected]> --- .../rocketmq/source/RocketMqSourceReader.java | 53 ++++- .../source/RocketMqConsumerThreadCloseTest.java | 217 +++++++++++++++++++++ 2 files changed, 267 insertions(+), 3 deletions(-) diff --git a/seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceReader.java b/seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceReader.java index bda2a49cc7..4072ae8498 100644 --- a/seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceReader.java +++ b/seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceReader.java @@ -42,6 +42,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; @Slf4j @@ -49,6 +50,16 @@ public class RocketMqSourceReader implements SourceReader<SeaTunnelRow, RocketMq private static final long THREAD_WAIT_TIME = 500L; + // Extra headroom added to a single poll timeout when bounding how long close() waits for + // every consumer thread to shut down concurrently. + private static final long CLOSE_GRACE_MILLIS = 500L; + + // Fallback when the configured poll timeout is unset (ConsumerMetadata#getBaseConfig()'s + // pollTimeoutMillis is a boxed Long with no default): close() still needs a concrete, finite + // bound to wait for concurrent consumer thread shutdown, so fall back to the connector's own + // documented default poll timeout instead of risking an NPE from unboxing a null value. + private static final long DEFAULT_POLL_TIMEOUT_MILLIS = 5000L; + private final Context context; private final ConsumerMetadata metadata; private final Set<RocketMqSourceSplit> sourceSplits; @@ -87,9 +98,45 @@ public class RocketMqSourceReader implements SourceReader<SeaTunnelRow, RocketMq @Override public void close() throws IOException { running = false; - consumerThreads.values().forEach(RocketMqConsumerThread::close); - if (executorService != null) { - executorService.shutdownNow(); + // DefaultLitePullConsumer's shutdown() is internally synchronized and can block for up + // to one poll cycle behind an in-flight poll(). Close every consumer thread concurrently + // instead of one at a time so the calling (task-cancellation) thread waits at most one + // poll timeout in total rather than pollTimeoutMillis * partitionCount, and isolate each + // thread's failure so one partition's shutdown exception can't abort the rest: a plain + // forEach here would stop at the first exception, leaking every remaining consumer + // thread and skipping executorService.shutdownNow() below entirely. + List<CompletableFuture<Void>> closeFutures = + consumerThreads.values().stream() + .map( + thread -> + CompletableFuture.runAsync( + () -> closeQuietly(thread), executorService)) + .collect(Collectors.toList()); + try { + CompletableFuture.allOf(closeFutures.toArray(new CompletableFuture[0])) + .get(pollTimeoutMillisOrDefault() + CLOSE_GRACE_MILLIS, TimeUnit.MILLISECONDS); + } catch (Exception e) { + log.warn("Timed out waiting for RocketMQ consumer threads to close", e); + } finally { + if (executorService != null) { + executorService.shutdownNow(); + } + } + } + + private long pollTimeoutMillisOrDefault() { + Long pollTimeoutMillis = + metadata.getBaseConfig() == null + ? null + : metadata.getBaseConfig().getPollTimeoutMillis(); + return pollTimeoutMillis == null ? DEFAULT_POLL_TIMEOUT_MILLIS : pollTimeoutMillis; + } + + private void closeQuietly(RocketMqConsumerThread thread) { + try { + thread.close(); + } catch (Exception e) { + log.warn("Failed to close a RocketMQ consumer thread", e); } } diff --git a/seatunnel-connectors-v2/connector-rocketmq/src/test/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqConsumerThreadCloseTest.java b/seatunnel-connectors-v2/connector-rocketmq/src/test/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqConsumerThreadCloseTest.java new file mode 100644 index 0000000000..7808cb0965 --- /dev/null +++ b/seatunnel-connectors-v2/connector-rocketmq/src/test/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqConsumerThreadCloseTest.java @@ -0,0 +1,217 @@ +/* + * 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.seatunnel.connectors.seatunnel.rocketmq.source; + +import org.apache.seatunnel.connectors.seatunnel.rocketmq.common.RocketMqAdminUtil; + +import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer; + +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +/** + * Covers the close-ordering fix for {@link RocketMqConsumerThread} and {@link + * RocketMqSourceReader}: a reader must stop the underlying RocketMQ client of every consumer thread + * it owns before it interrupts the thread pool. {@code DefaultLitePullConsumer#poll(long)} blocks + * on network I/O that plain thread interruption does not reliably unblock, so relying on {@code + * ExecutorService#shutdownNow()} alone could leave a consumer thread stuck mid-poll instead of + * returning promptly through {@link RocketMqConsumerThread#run()}'s {@code finally} block. + * + * <p>Also covers two follow-up hardening fixes to {@link RocketMqSourceReader#close()}: one + * consumer thread's shutdown failure must not abort closing the rest (a plain {@code forEach} would + * stop at the first exception), and every consumer thread must be closed concurrently rather than + * one at a time, since {@code DefaultLitePullConsumer}'s internally synchronized methods mean a + * sequential close could block the calling thread for {@code pollTimeoutMillis * partitionCount} + * instead of a single poll timeout. + */ +class RocketMqConsumerThreadCloseTest { + + /** + * {@link RocketMqConsumerThread#close()} must shut down its RocketMQ client directly, not only + * rely on the run-loop's own {@code finally} block, which only fires once the thread notices + * the interrupt - something a blocking {@code poll()} call is not guaranteed to do. + */ + @Test + void closeShutsDownUnderlyingConsumerClient() throws Exception { + DefaultLitePullConsumer mockConsumer = mock(DefaultLitePullConsumer.class); + RocketMqConsumerThread consumerThread; + try (MockedStatic<RocketMqAdminUtil> adminUtil = + Mockito.mockStatic(RocketMqAdminUtil.class)) { + adminUtil + .when(() -> RocketMqAdminUtil.initDefaultLitePullConsumer(any(), anyBoolean())) + .thenReturn(mockConsumer); + consumerThread = new RocketMqConsumerThread(new ConsumerMetadata()); + } + verify(mockConsumer, never()).shutdown(); + + consumerThread.close(); + + verify(mockConsumer, times(1)).shutdown(); + } + + /** + * {@link RocketMqSourceReader#close()} must mark the reader as no longer running and close + * every consumer thread it has created, in addition to shutting down its executor - dropping + * the explicit per-thread close would silently reintroduce the hang this fix removes. + */ + @Test + @SuppressWarnings("unchecked") + void closeStopsRunningAndClosesEveryTrackedConsumerThread() throws Exception { + RocketMqSourceReader reader = + new RocketMqSourceReader( + new ConsumerMetadata(), + new ConcurrentHashMap<>(), + mock(org.apache.seatunnel.api.source.SourceReader.Context.class)); + + RocketMqConsumerThread mockConsumerThread = mock(RocketMqConsumerThread.class); + setPrivateField(reader, "running", true); + Map<Object, RocketMqConsumerThread> consumerThreads = + (Map<Object, RocketMqConsumerThread>) getPrivateField(reader, "consumerThreads"); + consumerThreads.put(new Object(), mockConsumerThread); + + reader.close(); + + verify(mockConsumerThread, times(1)).close(); + assertFalse((boolean) getPrivateField(reader, "running")); + ExecutorService executorService = + (ExecutorService) getPrivateField(reader, "executorService"); + assertTrue(executorService.isShutdown()); + } + + /** + * One consumer thread's {@code close()} failure must not prevent the reader from closing every + * other tracked consumer thread or from shutting down its executor - a plain {@code forEach} + * would abort at the first exception and leak every remaining consumer thread plus skip {@code + * executorService.shutdownNow()} entirely. + */ + @Test + @SuppressWarnings("unchecked") + void closeIsolatesOneThreadsFailureFromTheRest() throws Exception { + RocketMqSourceReader reader = + new RocketMqSourceReader( + new ConsumerMetadata(), + new ConcurrentHashMap<>(), + mock(org.apache.seatunnel.api.source.SourceReader.Context.class)); + + RocketMqConsumerThread failingThread = mock(RocketMqConsumerThread.class); + RocketMqConsumerThread healthyThread = mock(RocketMqConsumerThread.class); + Mockito.doThrow(new RuntimeException("simulated shutdown failure")) + .when(failingThread) + .close(); + Map<Object, RocketMqConsumerThread> consumerThreads = + (Map<Object, RocketMqConsumerThread>) getPrivateField(reader, "consumerThreads"); + consumerThreads.put(new Object(), failingThread); + consumerThreads.put(new Object(), healthyThread); + + reader.close(); + + verify(failingThread, times(1)).close(); + verify(healthyThread, times(1)).close(); + ExecutorService executorService = + (ExecutorService) getPrivateField(reader, "executorService"); + assertTrue(executorService.isShutdown()); + } + + /** + * Consumer threads must be closed concurrently, not one at a time: {@code + * DefaultLitePullConsumer#shutdown()} is internally synchronized and can block behind an + * in-flight {@code poll()}, so closing sequentially would let the calling thread's wait grow + * with the partition count instead of staying bounded by a single poll timeout. Proven here by + * a rendezvous: if every mocked thread's {@code close()} is entered before any of them is + * allowed to return, they were not run one at a time. + */ + @Test + @SuppressWarnings("unchecked") + void closeShutsDownEveryConsumerThreadConcurrently() throws Exception { + int threadCount = 3; + RocketMqSourceReader reader = + new RocketMqSourceReader( + new ConsumerMetadata(), + new ConcurrentHashMap<>(), + mock(org.apache.seatunnel.api.source.SourceReader.Context.class)); + Map<Object, RocketMqConsumerThread> consumerThreads = + (Map<Object, RocketMqConsumerThread>) getPrivateField(reader, "consumerThreads"); + + CountDownLatch allEnteredClose = new CountDownLatch(threadCount); + CountDownLatch releaseClose = new CountDownLatch(1); + for (int i = 0; i < threadCount; i++) { + RocketMqConsumerThread thread = mock(RocketMqConsumerThread.class); + Mockito.doAnswer( + invocation -> { + allEnteredClose.countDown(); + releaseClose.await(5, TimeUnit.SECONDS); + return null; + }) + .when(thread) + .close(); + consumerThreads.put(new Object(), thread); + } + + Thread closer = + new Thread( + () -> { + try { + reader.close(); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + closer.start(); + try { + // If close() ran the mocked threads sequentially, only the first would ever enter + // close() - it would block on releaseClose forever, since nothing releases it yet, + // and this latch would time out instead of reaching zero. + assertTrue( + allEnteredClose.await(5, TimeUnit.SECONDS), + "expected every consumer thread to enter close() concurrently"); + } finally { + releaseClose.countDown(); + closer.join(TimeUnit.SECONDS.toMillis(5)); + } + } + + private static Object getPrivateField(Object target, String name) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.get(target); + } + + private static void setPrivateField(Object target, String name, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } +}
