diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcher.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcher.java index c83a8f95a0044..612e5deb5a05f 100644 --- a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcher.java +++ b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcher.java @@ -134,15 +134,21 @@ public void run() { // Put an empty synchronization batch to the element queue. // When this batch is recycled, all the records emitted earlier // must have already been processed. - elementsQueue.put( - fetcherId(), + final RecordsWithSplitIds synchronizationBatch = new RecordsBySplits(Collections.emptyMap(), Collections.emptySet()) { @Override public void recycle() { super.recycle(); recordsProcessedLatch.countDown(); } - }); + }; + // A pending wakeUp makes put() return without enqueueing, and it consumes the + // wakeUp flag on the way out. Retry, because a dropped batch is never recycled and + // the latch below would never be pulled. + boolean enqueued = false; + while (!enqueued) { + enqueued = elementsQueue.put(fetcherId(), synchronizationBatch); + } } } catch (Throwable t) { errorHandler.accept(t); diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherTest.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherTest.java index aac2a8100c67b..e8a52da058b32 100644 --- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherTest.java +++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherTest.java @@ -311,6 +311,60 @@ public void go() throws Exception { fetcherThread.sync(); } + @Test + void testShutdownSynchronizationBatchIsRetriedAfterWakeUp() throws Exception { + TestingSplitReader splitReader = new TestingSplitReader<>(); + // A capacity of one lets another producer hold the only slot, so the shutdown + // synchronization batch has to wait for the queue to drain. + FutureCompletingBlockingQueue> queue = + new FutureCompletingBlockingQueue<>(1); + final SplitFetcher fetcher = createFetcher(splitReader, queue); + + try { + assertThat(queue.put(1, finishedSplitFetch("other-split"))).isTrue(); + // An unconsumed wakeUp for the fetcher's own index makes its next put on a full queue + // return false without enqueueing anything. + queue.wakeUpPuttingThread(0); + fetcher.shutdown(true); + + // Spawn a new fetcher thread to go through the shutdown sequence. + CheckedThread fetcherThread = + new CheckedThread() { + @Override + public void go() { + fetcher.run(); + } + }; + fetcherThread.start(); + + // Wait until the fetcher thread blocks, either on putting the synchronization batch or + // on the shutdown latch. + waitUntil( + () -> fetcherThread.getState() == WAITING, + Duration.ofSeconds(30), + "The fetcher thread should have blocked while shutting down"); + assertThat(splitReader.isClosed()) + .as("The split reader should have not been closed.") + .isFalse(); + + // Free the slot held by the other producer. + queue.poll().recycle(); + + waitUntil( + () -> queue.size() == 1, + Duration.ofSeconds(30), + "The fetcher should have enqueued its shutdown synchronization batch"); + // Recycling it pulls the shutdown latch. + queue.poll().recycle(); + + fetcherThread.sync(); + assertThat(splitReader.isClosed()).isTrue(); + } finally { + // Releases the fetcher thread if the synchronization batch was dropped. + fetcher.shutdown(false); + } + } + // ------------------------------------------------------------------------ // testing utils // ------------------------------------------------------------------------