From 44b219316b5b795fe24dae9df7c901925ce6752b Mon Sep 17 00:00:00 2001
From: Martijn Visser <2989614+MartijnVisser@users.noreply.github.com>
Date: Mon, 14 Sep 2026 22:16:47 +0200
Subject: [PATCH] [FLINK-40660][connector-base] Retry the shutdown
 synchronization put in SplitFetcher

Retry the put that hands over the empty synchronization batch, because put() returns false without enqueueing when the queue is full and the fetcher has an unconsumed wakeUp flag, which drops the batch and leaves the fetcher blocked on recordsProcessedLatch with its SplitReader never closed.

Generated-by: Claude Code (Claude Opus 5)
---
 .../source/reader/fetcher/SplitFetcher.java   | 12 +++--
 .../reader/fetcher/SplitFetcherTest.java      | 54 +++++++++++++++++++
 2 files changed, 63 insertions(+), 3 deletions(-)

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 c83a8f95a00445..612e5deb5a05f3 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<E> synchronizationBatch =
                         new RecordsBySplits<E>(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 aac2a8100c67be..e8a52da058b327 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<Object, TestingSourceSplit> 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<RecordsWithSplitIds<Object>> queue =
+                new FutureCompletingBlockingQueue<>(1);
+        final SplitFetcher<Object, TestingSourceSplit> 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
     // ------------------------------------------------------------------------