[FLINK-40657][connector-base] Release per-fetcher state in FutureCompletingBlockingQueue

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcher.java

@@ -188,24 +188,29 @@ boolean runOnce() {
         }
 
         // execute the task outside of lock, so that it can be woken up
-        boolean taskFinished;
+        boolean taskFinished = false;
+        boolean taskRunCompleted = false;
         try {
             taskFinished = task.run();
+            taskRunCompleted = true;
         } catch (Exception e) {
             throw new RuntimeException(
                     String.format(
                             "SplitFetcher thread %d received unexpected exception while polling the records",
                             id),
                     e);
-        }
-
-        // re-acquire lock as all post-processing steps, need it
-        lock.lock();
-        try {
-            this.runningTask = null;
-            processTaskResultUnsafe(task, taskFinished);
         } finally {
-            lock.unlock();
+            // Re-acquire the lock because clearing the current task and processing a successful
+            // result must be atomic with respect to wakeup and shutdown.
+            lock.lock();
+            try {
+                this.runningTask = null;
+                if (taskRunCompleted) {
+                    processTaskResultUnsafe(task, taskFinished);
+                }
+            } finally {
+                lock.unlock();
+            }
         }
         return true;
     }

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManager.java

@@ -260,6 +260,7 @@ protected synchronized SplitFetcher<E, SplitT> createSplitFetcher() {
                         errorHandler,
                         () -> {
                             fetchers.remove(fetcherId);
+                            elementsQueue.releaseProducer(fetcherId);
                             fetchersToShutDown.decrementAndGet();
                             // We need this to synchronize status of fetchers to concurrent partners
                             // as

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueue.java

@@ -27,7 +27,8 @@
 
 import java.lang.reflect.Field;
 import java.util.ArrayDeque;
-import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -37,6 +38,7 @@
 import java.util.concurrent.locks.ReentrantLock;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkState;
 
 /**
  * A custom implementation of blocking queue in combination with a {@link CompletableFuture} that is
@@ -69,6 +71,11 @@
  * capacity limits, without interrupting the thread. This is done via the {@link
  * #wakeUpPuttingThread(int)} method.
  *
+ * <p>The queue keeps a small amount of per-producer state to support this. A producer that is
+ * permanently done with the queue must be handed to {@link #releaseProducer(int)}, the lifecycle
+ * counterpart of {@link #wakeUpPuttingThread(int)}, otherwise that state is retained for the
+ * lifetime of the queue.
+ *
  * @param <T> the type of the elements in the queue.
  */
 @Internal
@@ -102,9 +109,15 @@ public class FutureCompletingBlockingQueue<T> {
     @GuardedBy("lock")
     private final Queue<Condition> notFull;
 
-    /** The per-thread conditions and wakeUp flags. */
+    /**
+     * The per-producer conditions and wakeUp flags, keyed by the producer's thread index.
+     *
+     * <p>Entries are created on demand by {@link #conditionAndFlagFor(int)} and removed by {@link
+     * #releaseProducer(int)}, so this map is bounded by the peak number of live or in-flight
+     * producers rather than by the largest producer index ever allocated.
+     */
     @GuardedBy("lock")
-    private ConditionAndFlag[] putConditionAndFlags;
+    private final Map<Integer, ConditionAndFlag> putConditionAndFlags;
 
     public FutureCompletingBlockingQueue() {
         this(SourceReaderOptions.ELEMENT_QUEUE_CAPACITY.defaultValue());
@@ -115,7 +128,7 @@ public FutureCompletingBlockingQueue(int capacity) {
         this.capacity = capacity;
         this.queue = new ArrayDeque<>(capacity);
         this.lock = new ReentrantLock();
-        this.putConditionAndFlags = new ConditionAndFlag[1];
+        this.putConditionAndFlags = new HashMap<>();
         this.notFull = new ArrayDeque<>();
 
         // initially the queue is empty and thus unavailable
@@ -330,12 +343,55 @@ int getNumberOfQueuedPutters() {
     public void wakeUpPuttingThread(int threadIndex) {
         lock.lock();
         try {
-            maybeCreateCondition(threadIndex);
-            ConditionAndFlag caf = putConditionAndFlags[threadIndex];
-            if (caf != null) {
-                caf.setWakeUp(true);
-                caf.condition().signal();
+            // Creates the entry when absent, deliberately: the flag has to be sticky, so that a
+            // producer woken before it ever calls put() still observes the request and returns
+            // immediately instead of parking on a full queue.
+            final ConditionAndFlag caf = conditionAndFlagFor(threadIndex);
+            caf.setWakeUp(true);
+            caf.condition().signal();
+        } finally {
+            lock.unlock();
+        }
+    }
+
+    /**
+     * Releases the per-producer wakeup state held for {@code threadIndex}.
+     *
+     * <p>Call this once the producer with that index is permanently finished with this queue.
+     * Without it the queue retains each condition created by a wakeup request or by a put attempt
+     * made while the queue was full. {@code SplitFetcherManager} hands out a fresh, never-recycled
+     * index per {@code SplitFetcher}, so for a source whose fetchers are short-lived that set
+     * otherwise grows without bound for the lifetime of the JVM.
+     *
+     * <p>The call is idempotent and safe for an index that was never used.
+     *
+     * <p><b>The caller must not be inside {@link #put(int, Object)} for this index.</b> Releasing a
+     * producer that is still running would drop a pending wakeUp flag it has yet to observe, and it
+     * could then park on a full queue after having been told to stop. {@code SplitFetcher}
+     * satisfies this by running its shutdown hook only after its run loop has exited.
+     *
+     * <p>For the same reason this must be the last interaction with the queue for that index: a
+     * later {@link #wakeUpPuttingThread(int)} would recreate the entry, and nothing would remove it
+     * again. {@code SplitFetcher} satisfies this in its normal lifecycle because it clears the
+     * current task before its run loop exits and invokes the shutdown hook afterward.
+     *
+     * @param threadIndex The number identifying the producer thread, as passed to {@link #put(int,
+     *     Object)}.
+     */
+    public void releaseProducer(int threadIndex) {
+        lock.lock();
+        try {
+            final ConditionAndFlag caf = putConditionAndFlags.get(threadIndex);
+            if (caf == null) {
+                return;
+            }
+            if (caf.hasWaitingPutter()) {
+                // The condition may already have been removed from notFull after being signalled,
+                // while its putter is still waiting to reacquire the lock. Dropping the entry in
+                // that interval would discard a wakeUp flag the putter has yet to observe.
+                return;
             }
+            putConditionAndFlags.remove(threadIndex);
         } finally {
             lock.unlock();
         }
@@ -370,15 +426,17 @@ private T dequeue() {
 
     @GuardedBy("lock")
     private void waitOnPut(int fetcherIndex) throws InterruptedException {
-        maybeCreateCondition(fetcherIndex);
-        Condition cond = putConditionAndFlags[fetcherIndex].condition();
-        notFull.add(cond);
+        final ConditionAndFlag caf = conditionAndFlagFor(fetcherIndex);
+        final Condition cond = caf.condition();
+        caf.startWaiting();
         try {
+            notFull.add(cond);
             cond.await();
         } finally {
             // drop the condition once the thread stops waiting, so a later signalNextPutter()
             // does not signal a putter that is no longer waiting
             notFull.remove(cond);
+            caf.stopWaiting();
         }
     }
 
@@ -389,22 +447,24 @@ private void signalNextPutter() {
         }
     }
 
+    /**
+     * Returns the state for {@code threadIndex}, creating it when absent. Only {@link
+     * #wakeUpPuttingThread(int)} and {@link #waitOnPut(int)} call this, so a producer that is never
+     * woken and never blocks costs nothing.
+     */
     @GuardedBy("lock")
-    private void maybeCreateCondition(int threadIndex) {
-        if (putConditionAndFlags.length < threadIndex + 1) {
-            putConditionAndFlags = Arrays.copyOf(putConditionAndFlags, threadIndex + 1);
-        }
-
-        if (putConditionAndFlags[threadIndex] == null) {
-            putConditionAndFlags[threadIndex] = new ConditionAndFlag(lock.newCondition());
-        }
+    private ConditionAndFlag conditionAndFlagFor(int threadIndex) {
+        return putConditionAndFlags.computeIfAbsent(
+                threadIndex, ignored -> new ConditionAndFlag(lock.newCondition()));
     }
 
     @GuardedBy("lock")
     private boolean getAndResetWakeUpFlag(int threadIndex) {
-        maybeCreateCondition(threadIndex);
-        if (putConditionAndFlags[threadIndex].getWakeUp()) {
-            putConditionAndFlags[threadIndex].setWakeUp(false);
+        // Deliberately does not create the state: an absent entry cannot carry a wakeUp flag, and
+        // waitOnPut() creates it a moment later if this producer goes on to block.
+        final ConditionAndFlag caf = putConditionAndFlags.get(threadIndex);
+        if (caf != null && caf.getWakeUp()) {
+            caf.setWakeUp(false);
             return true;
         }
         return false;
@@ -415,6 +475,7 @@ private boolean getAndResetWakeUpFlag(int threadIndex) {
     private static class ConditionAndFlag {
         private final Condition cond;
         private boolean wakeUp;
+        private int waitingPutters;
 
         private ConditionAndFlag(Condition cond) {
             this.cond = cond;
@@ -429,6 +490,19 @@ private boolean getWakeUp() {
             return wakeUp;
         }
 
+        private void startWaiting() {
+            waitingPutters++;
+        }
+
+        private void stopWaiting() {
+            checkState(waitingPutters > 0, "stopWaiting() without a matching startWaiting()");
+            waitingPutters--;
+        }
+
+        private boolean hasWaitingPutter() {
+            return waitingPutters > 0;
+        }
+
         private void setWakeUp(boolean value) {
             wakeUp = value;
         }

flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManagerTest.java

@@ -32,6 +32,7 @@
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange;
 import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
+import org.apache.flink.connector.base.source.reader.synchronization.QueueProbe;
 import org.apache.flink.core.testutils.OneShotLatch;
 import org.apache.flink.util.MdcUtils;
 
@@ -140,6 +141,83 @@ void testCloseCleansUpPreviouslyClosedFetcher() throws Exception {
         fetcherManager.close(Long.MAX_VALUE);
     }
 
+    /**
+     * Each completed fetcher lifecycle must release its queue state, so monotonically increasing
+     * fetcher ids do not accumulate historical state.
+     */
+    @Test
+    @Timeout(value = 60, unit = TimeUnit.SECONDS)
+    void testFetcherShutdownReleasesWakeupStateAcrossLifecycles() throws Exception {
+        final String splitId = "testSplit";
+        final Configuration config = new Configuration();
+        config.set(SourceReaderOptions.ELEMENT_QUEUE_CAPACITY, 1);
+
+        final SplitFetcherManager<Integer, TestingSourceSplit> fetcherManager =
+                new SingleThreadFetcherManager<>(
+                        () ->
+                                new AwaitingReader<>(
+                                        new IOException("Should not happen"),
+                                        new RecordsBySplits<>(
+                                                Collections.emptyMap(),
+                                                Collections.singleton(splitId))),
+                        config);
+        final FutureCompletingBlockingQueue<RecordsWithSplitIds<Integer>> queue =
+                fetcherManager.getQueue();
+
+        try {
+            for (int expectedFetcherId = 0; expectedFetcherId < 3; expectedFetcherId++) {
+                fetcherManager.addSplits(
+                        Collections.singletonList(new TestingSourceSplit(splitId)));
+                assertThat(fetcherManager.fetchers).hasSize(1);
+                assertThat(fetcherManager.fetchers.keySet().iterator().next())
+                        .isEqualTo(expectedFetcherId);
+
+                waitUntil(
+                        () -> queue.size() == 1,
+                        Duration.ofSeconds(10),
+                        "The data batch should have filled the element queue.");
+                waitUntil(
+                        () -> {
+                            fetcherManager.maybeShutdownFinishedFetchers();
+                            return fetcherManager.fetchers.isEmpty();
+                        },
+                        Duration.ofSeconds(10),
+                        "The idle fetcher should have been removed.");
+                waitUntil(
+                        () -> QueueProbe.queuedPutters(queue) == 1,
+                        Duration.ofSeconds(10),
+                        "The final synchronization batch should be waiting for queue capacity.");
+
+                assertThat(QueueProbe.liveProducerStates(queue)).isOne();
+                assertThat(QueueProbe.producerStateStorageSize(queue)).isOne();
+
+                final RecordsWithSplitIds<Integer> dataBatch = queue.poll();
+                assertThat(dataBatch).isNotNull();
+                dataBatch.recycle();
+
+                waitUntil(
+                        () -> queue.size() == 1,
+                        Duration.ofSeconds(10),
+                        "The final synchronization batch should have been enqueued.");
+                final RecordsWithSplitIds<Integer> synchronizationBatch = queue.poll();
+                assertThat(synchronizationBatch).isNotNull();
+                synchronizationBatch.recycle();
+
+                waitUntil(
+                        () -> QueueProbe.liveProducerStates(queue) == 0,
+                        Duration.ofSeconds(10),
+                        "The shutdown hook should release the fetcher's queue state.");
+                assertThat(QueueProbe.producerStateStorageSize(queue)).isZero();
+            }
+        } finally {
+            RecordsWithSplitIds<Integer> batch;
+            while ((batch = queue.poll()) != null) {
+                batch.recycle();
+            }
+            fetcherManager.close(10_000L);
+        }
+    }
+
     /**
      * This test is somewhat testing the implementation instead of contract. This is because the
      * test is trying to make sure the element queue draining thread is not tight looping.