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..189e62b41d5ed 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 @@ -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; } diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManager.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManager.java index 6ea31d2d53bae..2f5fbe411553b 100644 --- a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManager.java +++ b/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 createSplitFetcher() { errorHandler, () -> { fetchers.remove(fetcherId); + elementsQueue.releaseProducer(fetcherId); fetchersToShutDown.decrementAndGet(); // We need this to synchronize status of fetchers to concurrent partners // as diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueue.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueue.java index 799fd0c4489b9..dc915e73f0d0d 100644 --- a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueue.java +++ b/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. * + *

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 the type of the elements in the queue. */ @Internal @@ -102,9 +109,15 @@ public class FutureCompletingBlockingQueue { @GuardedBy("lock") private final Queue notFull; - /** The per-thread conditions and wakeUp flags. */ + /** + * The per-producer conditions and wakeUp flags, keyed by the producer's thread index. + * + *

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 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}. + * + *

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. + * + *

The call is idempotent and safe for an index that was never used. + * + *

The caller must not be inside {@link #put(int, Object)} for this index. 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. + * + *

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; } diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManagerTest.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManagerTest.java index cf0fce5d42cf6..54e1ae29a3732 100644 --- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/fetcher/SplitFetcherManagerTest.java +++ b/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 fetcherManager = + new SingleThreadFetcherManager<>( + () -> + new AwaitingReader<>( + new IOException("Should not happen"), + new RecordsBySplits<>( + Collections.emptyMap(), + Collections.singleton(splitId))), + config); + final FutureCompletingBlockingQueue> 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 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 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 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. 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..9dd0da796665f 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 @@ -31,6 +31,7 @@ import org.junit.jupiter.api.Test; +import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; @@ -44,6 +45,7 @@ import static java.lang.Thread.State.WAITING; import static org.apache.flink.test.util.TestUtils.waitUntil; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Unit test for {@link SplitFetcher}. */ class SplitFetcherTest { @@ -259,6 +261,48 @@ void testClose() { assertThat(splitReader.isClosed()).isTrue(); } + @Test + void testRunningTaskIsClearedWhenTaskThrowsException() { + final SplitFetcher fetcher = + createFetcher(new TestingSplitReader<>()); + final IOException failure = new IOException("test failure"); + final FailingFetcherTask task = new FailingFetcherTask(failure); + fetcher.enqueueTask(task); + + try { + assertThatThrownBy(fetcher::runOnce) + .isInstanceOf(RuntimeException.class) + .hasCause(failure); + + assertThat(fetcher.isIdle()).isTrue(); + fetcher.wakeUp(false); + assertThat(task.getNumWakeUps()).isZero(); + } finally { + fetcher.shutdown(); + fetcher.run(); + } + } + + @Test + void testRunningTaskIsClearedWhenTaskThrowsError() { + final SplitFetcher fetcher = + createFetcher(new TestingSplitReader<>()); + final AssertionError failure = new AssertionError("test failure"); + final FailingFetcherTask task = new FailingFetcherTask(failure); + fetcher.enqueueTask(task); + + try { + assertThatThrownBy(fetcher::runOnce).isSameAs(failure); + + assertThat(fetcher.isIdle()).isTrue(); + fetcher.wakeUp(false); + assertThat(task.getNumWakeUps()).isZero(); + } finally { + fetcher.shutdown(); + fetcher.run(); + } + } + @Test void testCloseAfterPause() throws InterruptedException { final FutureCompletingBlockingQueue> queue = @@ -319,6 +363,33 @@ private static RecordsBySplits finishedSplitFetch(String splitId) { return new RecordsBySplits<>(Collections.emptyMap(), Collections.singleton(splitId)); } + private static final class FailingFetcherTask implements SplitFetcherTask { + + private final Throwable failure; + private final AtomicInteger numWakeUps = new AtomicInteger(); + + private FailingFetcherTask(Throwable failure) { + this.failure = failure; + } + + @Override + public boolean run() throws IOException { + if (failure instanceof IOException) { + throw (IOException) failure; + } + throw (Error) failure; + } + + @Override + public void wakeUp() { + numWakeUps.incrementAndGet(); + } + + private int getNumWakeUps() { + return numWakeUps.get(); + } + } + private static SplitFetcher createFetcher( final SplitReader reader) { return createFetcher(reader, new FutureCompletingBlockingQueue<>()); diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueueTest.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueueTest.java index 65e7ff38b9239..68dc92dd8c3d8 100644 --- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueueTest.java +++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/FutureCompletingBlockingQueueTest.java @@ -33,6 +33,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; @@ -86,24 +87,11 @@ void testPollEmptyQueue() throws InterruptedException { @Test void testWakeUpPut() throws InterruptedException { - FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); - - CountDownLatch latch = new CountDownLatch(1); - new Thread( - () -> { - try { - assertThat(queue.put(0, 1234)).isTrue(); - assertThat(queue.put(0, 1234)).isFalse(); - latch.countDown(); - } catch (InterruptedException e) { - fail("Interrupted unexpectedly."); - } - }) - .start(); + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); queue.wakeUpPuttingThread(0); - latch.await(); - assertThat(latch.getCount()).isEqualTo(0); + assertThat(queue.put(1, 1234)).isTrue(); + assertThat(queue.put(0, 1234)).isFalse(); } /** @@ -147,6 +135,233 @@ void testWakeUpDoesNotStrandAnotherPutter() throws Exception { .isTrue(); } + /** + * Without {@link FutureCompletingBlockingQueue#releaseProducer(int)} the queue keeps one + * condition per producer index it has ever seen. Because {@code SplitFetcherManager} allocates + * a fresh, never-recycled index per {@code SplitFetcher}, a source with short-lived fetchers + * accumulates them for the lifetime of the JVM. + */ + @Test + void testReleaseProducerBoundsTheWakeupState() throws InterruptedException { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + + for (int producer = 0; producer < 3; producer++) { + // Fill the single slot, so the next put takes the full-queue path that registers + // wakeup state for this producer. + assertThat(queue.put(producer, producer)).isTrue(); + queue.wakeUpPuttingThread(producer); + assertThat(queue.put(producer, producer)).isFalse(); + queue.poll(); + + assertThat(QueueProbe.liveProducerStates(queue)).isOne(); + assertThat(QueueProbe.producerStateStorageSize(queue)).isOne(); + + queue.releaseProducer(producer); + + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + assertThat(QueueProbe.producerStateStorageSize(queue)).isZero(); + } + + assertThat(queue.getNumberOfQueuedPutters()).isZero(); + } + + /** + * The storage assertions are what separate releasing the state from merely clearing it. An + * implementation that kept the {@code ConditionAndFlag[]} and only nulled the released slot + * would satisfy every live-count assertion above while still growing its backing array to the + * largest index ever seen, so this pins the storage down as well. + */ + @Test + void testSparseProducerIndexDoesNotExpandStorage() throws InterruptedException { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + + assertThat(queue.put(100_000, 1)).isTrue(); + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + assertThat(QueueProbe.producerStateStorageSize(queue)).isZero(); + assertThat(queue.poll()).isOne(); + + queue.wakeUpPuttingThread(100_000); + + assertThat(QueueProbe.liveProducerStates(queue)).isOne(); + assertThat(QueueProbe.producerStateStorageSize(queue)).isOne(); + + queue.releaseProducer(100_000); + + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + assertThat(QueueProbe.producerStateStorageSize(queue)).isZero(); + } + + /** Release is idempotent, tolerates unknown ids, and only removes the target state. */ + @Test + void testReleaseProducerOnlyRemovesTheTargetState() { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + + queue.wakeUpPuttingThread(3); + queue.wakeUpPuttingThread(4); + + queue.releaseProducer(7); + queue.releaseProducer(3); + queue.releaseProducer(3); + + assertThat(QueueProbe.containsProducerState(queue, 3)).isFalse(); + assertThat(QueueProbe.containsProducerState(queue, 4)).isTrue(); + assertThat(QueueProbe.liveProducerStates(queue)).isOne(); + + queue.releaseProducer(4); + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + } + + /** + * Releasing a producer that is currently parked in {@code waitOnPut} would discard the wakeUp + * flag it is about to read and could leave it parked for good, so the release must be refused + * while it is still waiting. + */ + @Test + @Timeout(value = 60, unit = TimeUnit.SECONDS) + void testReleaseProducerIsRefusedWhileTheProducerIsParked() throws Exception { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + + queue.put(0, 0); + + final AtomicBoolean parkedPutterResult = new AtomicBoolean(true); + final Thread parkedPutter = + new Thread(() -> parkedPutterResult.set(putUnchecked(queue, 1, 1)), "parkedPutter"); + parkedPutter.start(); + try { + CommonTestUtils.waitUntilCondition(() -> queue.getNumberOfQueuedPutters() == 1); + + queue.releaseProducer(1); + assertThat(QueueProbe.liveProducerStates(queue)) + .as("must not drop wakeup state for a producer currently inside put()") + .isOne(); + + // The graceful wakeup still reaches it, which is what the refusal protects. + queue.wakeUpPuttingThread(1); + joinWithinTimeout(parkedPutter); + assertThat(parkedPutterResult).isFalse(); + + // Once it has left put(), the release goes through. + queue.releaseProducer(1); + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + } finally { + if (parkedPutter.isAlive()) { + queue.wakeUpPuttingThread(1); + parkedPutter.interrupt(); + queue.poll(); + parkedPutter.join(TimeUnit.SECONDS.toMillis(10)); + } + queue.releaseProducer(1); + } + } + + /** An interrupted wait must not leave the producer permanently marked as waiting. */ + @Test + @Timeout(value = 30, unit = TimeUnit.SECONDS) + void testInterruptedPutterDoesNotPreventProducerRelease() throws Exception { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + assertThat(queue.put(0, 0)).isTrue(); + + final CompletableFuture putInterrupted = new CompletableFuture<>(); + final Thread putter = + new Thread( + () -> { + try { + queue.put(1, 1); + putInterrupted.complete(false); + } catch (InterruptedException expected) { + putInterrupted.complete(true); + } catch (Throwable failure) { + putInterrupted.completeExceptionally(failure); + } + }, + "interruptedPutter"); + putter.start(); + try { + CommonTestUtils.waitUntilCondition(() -> queue.getNumberOfQueuedPutters() == 1); + + putter.interrupt(); + assertThat(putInterrupted.get(10, TimeUnit.SECONDS)).isTrue(); + joinWithinTimeout(putter); + assertThat(queue.getNumberOfQueuedPutters()).isZero(); + assertThat(QueueProbe.containsProducerState(queue, 1)).isTrue(); + + queue.releaseProducer(1); + assertThat(QueueProbe.containsProducerState(queue, 1)).isFalse(); + } finally { + if (putter.isAlive()) { + putter.interrupt(); + queue.poll(); + putter.join(TimeUnit.SECONDS.toMillis(10)); + } + queue.releaseProducer(1); + } + } + + /** + * The companion to the test above, for the window that makes membership of {@code notFull} an + * unreliable answer to "is this producer still inside {@code put()}?". + * + *

{@code signalNextPutter()} removes a condition from {@code notFull} at signal time, not + * when its producer resumes, so between the signal and that producer reacquiring the lock it is + * still inside {@code put()} while absent from {@code notFull}. A release in that window would + * drop the state together with a wakeUp flag the producer has not read yet, and the producer + * would then build fresh state with no flag and park again. The {@code waitingPutters} counter + * spans the whole of {@code cond.await()} and therefore covers it. + * + *

The window is driven deterministically rather than raced for: the test thread takes the + * queue's own lock, so the signalled producer cannot reacquire it and cannot leave {@code + * put()} until the test releases it. + */ + @Test + @Timeout(value = 30, unit = TimeUnit.SECONDS) + void testReleaseProducerIsRefusedAfterSignalBeforePutterReacquiresLock() throws Exception { + final FutureCompletingBlockingQueue queue = new FutureCompletingBlockingQueue<>(1); + assertThat(queue.put(0, 0)).isTrue(); + + final AtomicBoolean putterResult = new AtomicBoolean(true); + final Thread putter = + new Thread(() -> putterResult.set(putUnchecked(queue, 1, 1)), "signalledPutter"); + putter.start(); + try { + CommonTestUtils.waitUntilCondition(() -> queue.getNumberOfQueuedPutters() == 1); + + final ReentrantLock queueLock = QueueProbe.queueLock(queue); + queueLock.lock(); + try { + assertThat(queue.poll()).isZero(); + assertThat(QueueProbe.queuedPutters(queue)).isZero(); + + queue.wakeUpPuttingThread(1); + queue.releaseProducer(1); + assertThat(QueueProbe.containsProducerState(queue, 1)).isTrue(); + + assertThat(queue.put(2, 2)).isTrue(); + } finally { + queueLock.unlock(); + } + + joinWithinTimeout(putter); + assertThat(putterResult).isFalse(); + + queue.releaseProducer(1); + assertThat(QueueProbe.liveProducerStates(queue)).isZero(); + } finally { + if (putter.isAlive()) { + queue.wakeUpPuttingThread(1); + putter.interrupt(); + queue.poll(); + putter.join(TimeUnit.SECONDS.toMillis(10)); + } + queue.releaseProducer(1); + queue.poll(); + } + } + + private static void joinWithinTimeout(Thread thread) throws InterruptedException { + thread.join(TimeUnit.SECONDS.toMillis(10)); + assertThat(thread.isAlive()).as("The putting thread should have terminated").isFalse(); + } + private static boolean putUnchecked( FutureCompletingBlockingQueue queue, int threadIndex, int value) { try { diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/QueueProbe.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/QueueProbe.java new file mode 100644 index 0000000000000..bc59d7e918be8 --- /dev/null +++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/synchronization/QueueProbe.java @@ -0,0 +1,148 @@ +/* + * 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.flink.connector.base.source.reader.synchronization; + +import java.lang.reflect.Array; +import java.lang.reflect.Field; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; + +/** + * Test-only access to the internal state of {@link FutureCompletingBlockingQueue}. + * + *

Reflection is used deliberately rather than adding accessors to the queue: a {@code + * VisibleForTesting} method in connector production code registers as a dependency on non-public + * Flink API and would need a new entry in the frozen architecture-violation store. + * + *

The producer-state probes handle both a map and an array representation so that they stay + * meaningful against an implementation that releases state by nulling an array slot. Such an + * implementation bounds {@link #liveProducerStates} but not {@link #producerStateStorageSize}, + * which is why both are exposed. + */ +public final class QueueProbe { + + private static final Field PRODUCER_STATES_FIELD = field("putConditionAndFlags"); + private static final Field NOT_FULL_FIELD = field("notFull"); + private static final Field LOCK_FIELD = field("lock"); + + private QueueProbe() {} + + /** Returns the number of producers the queue currently holds state for. */ + public static int liveProducerStates(FutureCompletingBlockingQueue queue) { + return withQueueLock( + queue, + () -> { + final Object states = producerStates(queue); + if (states instanceof Map) { + return ((Map) states).size(); + } + + int liveStates = 0; + for (int index = 0; index < Array.getLength(states); index++) { + if (Array.get(states, index) != null) { + liveStates++; + } + } + return liveStates; + }); + } + + /** + * Returns the size of the backing storage: map entries for a map implementation, allocated + * slots for an array. Unlike {@link #liveProducerStates} this does not fall back to zero when + * an array implementation nulls its released slots. + */ + public static int producerStateStorageSize(FutureCompletingBlockingQueue queue) { + return withQueueLock( + queue, + () -> { + final Object states = producerStates(queue); + if (states instanceof Map) { + return ((Map) states).size(); + } + return Array.getLength(states); + }); + } + + /** Returns whether the queue holds state for {@code producerIndex} specifically. */ + public static boolean containsProducerState( + FutureCompletingBlockingQueue queue, int producerIndex) { + return withQueueLock( + queue, + () -> { + final Object states = producerStates(queue); + if (states instanceof Map) { + return ((Map) states).containsKey(producerIndex); + } + return producerIndex >= 0 + && producerIndex < Array.getLength(states) + && Array.get(states, producerIndex) != null; + }); + } + + /** + * Returns the number of conditions in the queue's {@code notFull} waiter set. Note this drops a + * waiter as soon as it is signalled, before it has left {@code put()}. + */ + public static int queuedPutters(FutureCompletingBlockingQueue queue) { + return withQueueLock(queue, () -> ((Queue) read(NOT_FULL_FIELD, queue)).size()); + } + + /** + * Returns the queue's own lock, so a test can hold it to drive an interleaving + * deterministically instead of racing for it. + */ + public static ReentrantLock queueLock(FutureCompletingBlockingQueue queue) { + return (ReentrantLock) read(LOCK_FIELD, queue); + } + + private static Object producerStates(FutureCompletingBlockingQueue queue) { + return read(PRODUCER_STATES_FIELD, queue); + } + + private static T withQueueLock(FutureCompletingBlockingQueue queue, Supplier action) { + final ReentrantLock lock = queueLock(queue); + lock.lock(); + try { + return action.get(); + } finally { + lock.unlock(); + } + } + + private static Field field(String name) { + try { + final Field field = FutureCompletingBlockingQueue.class.getDeclaredField(name); + field.setAccessible(true); + return field; + } catch (NoSuchFieldException e) { + throw new ExceptionInInitializerError(e); + } + } + + private static Object read(Field field, FutureCompletingBlockingQueue queue) { + try { + return field.get(queue); + } catch (IllegalAccessException e) { + throw new AssertionError(e); + } + } +}