This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit da4a182934fe974f7bc821e931d6d2a7bf862e19 Author: luogen.lg <[email protected]> AuthorDate: Wed Jun 3 16:09:35 2026 +0800 [FLINK-39835][streaming-java] AsyncWaitOperator implements SupportsChainAvailability --- .../api/operators/async/AsyncWaitOperator.java | 36 +++++- .../async/queue/OrderedStreamElementQueue.java | 23 +++- .../operators/async/queue/StreamElementQueue.java | 3 +- .../async/queue/UnorderedStreamElementQueue.java | 23 +++- .../api/operators/async/AsyncWaitOperatorTest.java | 130 +++++++++++++++++++++ .../async/queue/OrderedStreamElementQueueTest.java | 51 ++++++++ .../queue/UnorderedStreamElementQueueTest.java | 52 +++++++++ .../tasks/StreamTaskMailboxTestHarness.java | 4 + 8 files changed, 318 insertions(+), 4 deletions(-) diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperator.java index 323fe4106f3..873e7c3ff67 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperator.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperator.java @@ -23,6 +23,7 @@ import org.apache.flink.api.common.operators.MailboxExecutor; import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.runtime.io.AvailabilityProvider; import org.apache.flink.runtime.state.StateInitializationContext; import org.apache.flink.runtime.state.StateSnapshotContext; import org.apache.flink.streaming.api.datastream.AsyncDataStream; @@ -37,6 +38,7 @@ import org.apache.flink.streaming.api.operators.BoundedOneInput; import org.apache.flink.streaming.api.operators.OneInputStreamOperator; import org.apache.flink.streaming.api.operators.Output; import org.apache.flink.streaming.api.operators.StreamOperatorParameters; +import org.apache.flink.streaming.api.operators.SupportsChainAvailability; import org.apache.flink.streaming.api.operators.TimestampedCollector; import org.apache.flink.streaming.api.operators.async.queue.OrderedStreamElementQueue; import org.apache.flink.streaming.api.operators.async.queue.StreamElementQueue; @@ -57,6 +59,7 @@ import java.util.Collections; import java.util.HashSet; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.atomic.AtomicBoolean; @@ -92,7 +95,7 @@ import static org.apache.flink.streaming.util.retryable.AsyncRetryStrategies.NO_ @Internal public class AsyncWaitOperator<IN, OUT> extends AbstractUdfStreamOperator<OUT, AsyncFunction<IN, OUT>> - implements OneInputStreamOperator<IN, OUT>, BoundedOneInput { + implements OneInputStreamOperator<IN, OUT>, BoundedOneInput, SupportsChainAvailability { private static final long serialVersionUID = 1L; private static final String STATE_NAME = "_async_wait_operator_state_"; @@ -139,6 +142,8 @@ public class AsyncWaitOperator<IN, OUT> /** Whether retry is disabled due to task finish, initially set to false. */ private transient AtomicBoolean retryDisabledOnFinish; + private AvailabilityProvider downstreamAvailabilityProvider; + public AsyncWaitOperator( StreamOperatorParameters<OUT> parameters, @Nonnull AsyncFunction<IN, OUT> asyncFunction, @@ -390,6 +395,25 @@ public class AsyncWaitOperator<IN, OUT> */ private void outputCompletedElement() { if (queue.hasCompletedElements()) { + // Defer emission until downstream becomes available. + if (downstreamAvailabilityProvider != null + && !downstreamAvailabilityProvider.isAvailable()) { + downstreamAvailabilityProvider + .getAvailableFuture() + .thenRun( + () -> { + try { + mailboxExecutor.execute( + this::outputCompletedElement, + "AsyncWaitOperator#outputCompletedElement(deferred)"); + } catch (RejectedExecutionException e) { + LOG.debug( + "Deferred element emission is ignored since the mailbox rejected the execution.", + e); + } + }); + return; + } // emit only one element to not block the mailbox thread unnecessarily queue.emitCompletedElement(timestampedCollector); // if there are more completed elements, emit them with subsequent mails @@ -431,6 +455,16 @@ public class AsyncWaitOperator<IN, OUT> timeoutTimestamp, timestamp -> callback.accept(null)); } + @Override + public CompletableFuture<?> getAvailableFuture() { + return queue.getAvailableFuture(); + } + + @Override + public void setDownstreamAvailabilityProvider(AvailabilityProvider provider) { + this.downstreamAvailabilityProvider = provider; + } + /** A delegator holds the real {@link ResultHandler} to handle retries. */ private class RetryableResultHandlerDelegator implements ResultFuture<OUT> { diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueue.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueue.java index 6d8fa95f94f..84a94a531d3 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueue.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueue.java @@ -34,6 +34,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Optional; import java.util.Queue; +import java.util.concurrent.CompletableFuture; /** * Ordered {@link StreamElementQueue} implementation. The ordered stream element queue provides @@ -52,11 +53,14 @@ public final class OrderedStreamElementQueue<OUT> implements StreamElementQueue< /** Queue for the inserted StreamElementQueueEntries. */ private final Queue<StreamElementQueueEntry<OUT>> queue; + private final AvailabilityHelper availabilityHelper = new AvailabilityHelper(); + public OrderedStreamElementQueue(int capacity) { Preconditions.checkArgument(capacity > 0, "The capacity must be larger than 0."); this.capacity = capacity; this.queue = new ArrayDeque<>(capacity); + this.availabilityHelper.resetAvailable(); } @Override @@ -67,8 +71,12 @@ public final class OrderedStreamElementQueue<OUT> implements StreamElementQueue< @Override public void emitCompletedElement(TimestampedCollector<OUT> output) { if (hasCompletedElements()) { + boolean isFullBefore = isFull(); final StreamElementQueueEntry<OUT> head = queue.poll(); head.emitResult(output); + if (isFullBefore && !isFull()) { + availabilityHelper.getUnavailableToResetAvailable().complete(null); + } } } @@ -93,7 +101,7 @@ public final class OrderedStreamElementQueue<OUT> implements StreamElementQueue< @Override public Optional<ResultFuture<OUT>> tryPut(StreamElement streamElement) { - if (queue.size() < capacity) { + if (!isFull()) { StreamElementQueueEntry<OUT> queueEntry = createEntry(streamElement); queue.add(queueEntry); @@ -104,6 +112,10 @@ public final class OrderedStreamElementQueue<OUT> implements StreamElementQueue< queue.size(), capacity); + if (isFull()) { + availabilityHelper.resetUnavailable(); + } + return Optional.of(queueEntry); } else { LOG.debug( @@ -116,6 +128,15 @@ public final class OrderedStreamElementQueue<OUT> implements StreamElementQueue< } } + @Override + public CompletableFuture<?> getAvailableFuture() { + return availabilityHelper.getAvailableFuture(); + } + + private boolean isFull() { + return queue.size() >= capacity; + } + private StreamElementQueueEntry<OUT> createEntry(StreamElement streamElement) { if (streamElement.isRecord()) { return new StreamRecordQueueEntry<>((StreamRecord<?>) streamElement); diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/StreamElementQueue.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/StreamElementQueue.java index bb6c18dc310..e26ab02959f 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/StreamElementQueue.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/StreamElementQueue.java @@ -19,6 +19,7 @@ package org.apache.flink.streaming.api.operators.async.queue; import org.apache.flink.annotation.Internal; +import org.apache.flink.runtime.io.AvailabilityProvider; import org.apache.flink.streaming.api.functions.async.ResultFuture; import org.apache.flink.streaming.api.operators.TimestampedCollector; import org.apache.flink.streaming.api.operators.async.AsyncWaitOperator; @@ -29,7 +30,7 @@ import java.util.Optional; /** Interface for stream element queues for the {@link AsyncWaitOperator}. */ @Internal -public interface StreamElementQueue<OUT> { +public interface StreamElementQueue<OUT> extends AvailabilityProvider { /** * Tries to put the given element in the queue. This operation succeeds if the queue has diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueue.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueue.java index 289b1b89087..7e79c2cb8cd 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueue.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueue.java @@ -38,6 +38,7 @@ import java.util.List; import java.util.Optional; import java.util.Queue; import java.util.Set; +import java.util.concurrent.CompletableFuture; /** * Unordered implementation of the {@link StreamElementQueue}. The unordered stream element queue @@ -63,6 +64,8 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu private int numberOfEntries; + private final AvailabilityHelper availabilityHelper = new AvailabilityHelper(); + public UnorderedStreamElementQueue(int capacity) { Preconditions.checkArgument(capacity > 0, "The capacity must be larger than 0."); @@ -70,11 +73,12 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu // most likely scenario are 4 segments <elements, watermark, elements, watermark> this.segments = new ArrayDeque<>(4); this.numberOfEntries = 0; + this.availabilityHelper.resetAvailable(); } @Override public Optional<ResultFuture<OUT>> tryPut(StreamElement streamElement) { - if (size() < capacity) { + if (!isFull()) { StreamElementQueueEntry<OUT> queueEntry; if (streamElement.isRecord()) { queueEntry = addRecord((StreamRecord<?>) streamElement); @@ -92,6 +96,10 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu size(), capacity); + if (isFull()) { + availabilityHelper.resetUnavailable(); + } + return Optional.of(queueEntry); } else { LOG.debug( @@ -155,6 +163,7 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu if (segments.isEmpty()) { return; } + boolean isFullBefore = isFull(); final Segment currentSegment = segments.getFirst(); numberOfEntries -= currentSegment.emitCompleted(output); @@ -163,6 +172,9 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu if (segments.size() > 1 && currentSegment.isEmpty()) { segments.pop(); } + if (isFullBefore && !isFull()) { + availabilityHelper.getUnavailableToResetAvailable().complete(null); + } } @Override @@ -184,6 +196,15 @@ public final class UnorderedStreamElementQueue<OUT> implements StreamElementQueu return numberOfEntries; } + @Override + public CompletableFuture<?> getAvailableFuture() { + return availabilityHelper.getAvailableFuture(); + } + + private boolean isFull() { + return size() >= capacity; + } + /** An entry that notifies the respective segment upon completion. */ static class SegmentedStreamRecordQueueEntry<OUT> extends StreamRecordQueueEntry<OUT> { private final Segment<OUT> segment; diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperatorTest.java index 6ed1fa364ea..2c8dc6fb880 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperatorTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperatorTest.java @@ -31,6 +31,7 @@ import org.apache.flink.runtime.checkpoint.CheckpointMetaData; import org.apache.flink.runtime.checkpoint.CheckpointOptions; import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; import org.apache.flink.runtime.checkpoint.TaskStateSnapshot; +import org.apache.flink.runtime.io.AvailabilityProvider; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; import org.apache.flink.runtime.jobgraph.JobGraph; import org.apache.flink.runtime.jobgraph.JobVertex; @@ -55,6 +56,7 @@ import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask; import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness; import org.apache.flink.streaming.runtime.tasks.StreamTaskMailboxTestHarness; import org.apache.flink.streaming.runtime.tasks.StreamTaskMailboxTestHarnessBuilder; +import org.apache.flink.streaming.runtime.tasks.mailbox.Mail; import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; import org.apache.flink.streaming.util.TestHarnessUtil; import org.apache.flink.streaming.util.retryable.AsyncRetryStrategies; @@ -486,6 +488,134 @@ public class AsyncWaitOperatorTest { } } + /** + * Tests that the OperatorChain availability provider exposed to {@link + * org.apache.flink.streaming.runtime.tasks.StreamTask} flips to unavailable when an {@link + * AsyncWaitOperator}'s queue is full, and recovers once a slot is freed by an emitted result. + * This is the integration that drives the chain-availability path in {@code + * StreamTask#processInput}. + */ + @Test + void testStreamTaskChainAvailabilityFlipsWithAsyncQueueCapacity() throws Exception { + SharedReference<List<ResultFuture<?>>> resultFutures = sharedObjects.add(new ArrayList<>()); + + StreamTaskMailboxTestHarnessBuilder<Integer> builder = + new StreamTaskMailboxTestHarnessBuilder<>( + OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO) + .addInput(BasicTypeInfo.INT_TYPE_INFO); + + try (StreamTaskMailboxTestHarness<Integer> harness = + builder.setupOutputForSingletonOperatorChain( + new AsyncWaitOperatorFactory<>( + new CollectableFuturesAsyncFunction<>(resultFutures), + TIMEOUT, + 2, + AsyncDataStream.OutputMode.ORDERED)) + .build()) { + + AvailabilityProvider chainAvailability = + harness.getOperatorChain().getChainAvailabilityProvider(); + assertThat(chainAvailability) + .as("Singleton AsyncWait chain must expose an availability provider") + .isNotNull(); + assertThat(chainAvailability.isAvailable()) + .as("Empty async queue starts available") + .isTrue(); + + // Fill the capacity-2 queue. The async function never auto-completes. + harness.processElement(new StreamRecord<>(1, 0L)); + harness.processElement(new StreamRecord<>(2, 1L)); + + assertThat(chainAvailability.isAvailable()) + .as( + "After filling capacity, chain availability must drop to false so" + + " StreamTask#processInput enters the chain-availability" + + " branch") + .isFalse(); + + // Complete one entry; the head emits and frees a slot, restoring availability. + completeWithSingle(resultFutures.get().get(0), 10); + harness.processAll(); + + assertThat(chainAvailability.isAvailable()) + .as("Freeing one slot must restore chain availability") + .isTrue(); + + // Drain the rest to keep the harness clean. + completeWithSingle(resultFutures.get().get(1), 20); + harness.processAll(); + assertThat(harness.getOutput()) + .containsExactly(new StreamRecord<>(10, 0L), new StreamRecord<>(20, 1L)); + } + } + + /** + * Tests that {@link AsyncWaitOperator} defers element emission while the configured downstream + * {@link AvailabilityProvider} is unavailable, and that completing the downstream's future + * schedules a deferred mail that finally drains the buffered element. + */ + @Test + void testOutputCompletedElementDefersWhenDownstreamUnavailable() throws Exception { + final OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = + createTestHarness( + new ImmediateAsyncFunction(), + TIMEOUT, + 2, + AsyncDataStream.OutputMode.ORDERED); + + testHarness.open(); + + final AsyncWaitOperator<Integer, Integer> operator = + (AsyncWaitOperator<Integer, Integer>) testHarness.getOneInputOperator(); + final CompletableFuture<Void> downstreamFuture = new CompletableFuture<>(); + operator.setDownstreamAvailabilityProvider(() -> downstreamFuture); + + synchronized (testHarness.getCheckpointLock()) { + testHarness.processElement(new StreamRecord<>(7, 0L)); + } + + // ImmediateAsyncFunction completes synchronously; draining runs the resulting + // outputCompletedElement mail, which must observe the unavailable downstream + // and register a thenRun callback instead of emitting. + drainMailbox(testHarness); + assertThat(testHarness.getOutput()) + .as("Element must not be emitted while downstream is unavailable") + .isEmpty(); + + // Signal downstream availability. The thenRun callback runs synchronously on + // this thread and enqueues a deferred outputCompletedElement mail. + downstreamFuture.complete(null); + drainMailbox(testHarness); + assertThat(testHarness.getOutput()).containsExactly(new StreamRecord<>(7, 0L)); + + synchronized (testHarness.getCheckpointLock()) { + testHarness.endInput(); + testHarness.close(); + } + } + + private static void drainMailbox(OneInputStreamOperatorTestHarness<?, ?> testHarness) + throws Exception { + for (Mail mail : testHarness.getTaskMailbox().drain()) { + mail.run(); + } + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static void completeWithSingle(ResultFuture<?> resultFuture, Object value) { + ((ResultFuture) resultFuture).complete(Collections.singletonList(value)); + } + + /** A synchronous {@link AsyncFunction} that completes the result on the calling thread. */ + private static class ImmediateAsyncFunction implements AsyncFunction<Integer, Integer> { + private static final long serialVersionUID = 1L; + + @Override + public void asyncInvoke(Integer input, ResultFuture<Integer> resultFuture) { + resultFuture.complete(Collections.singletonList(input)); + } + } + /** Tests that the AsyncWaitOperator works together with chaining. */ @Test void testOperatorChainWithProcessingTime() throws Exception { diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueueTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueueTest.java index e5f83ab901d..b79d094a1e7 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueueTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/OrderedStreamElementQueueTest.java @@ -28,6 +28,7 @@ import org.junit.jupiter.api.Test; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.CompletableFuture; import static org.apache.flink.streaming.api.operators.async.queue.QueueUtil.popCompleted; import static org.apache.flink.streaming.api.operators.async.queue.QueueUtil.putSuccessfully; @@ -70,4 +71,54 @@ public class OrderedStreamElementQueueTest { assertThat(queue.size()).isZero(); assertThat(queue.isEmpty()).isTrue(); } + + /** + * Tests that {@link OrderedStreamElementQueue#isAvailable()} flips to false once the queue is + * full and that {@link OrderedStreamElementQueue#getAvailableFuture()} returns an incomplete + * future in that state. + */ + @Test + void testIsAvailableTogglesWithCapacity() { + final OrderedStreamElementQueue<Integer> queue = new OrderedStreamElementQueue<>(2); + + assertThat(queue.isAvailable()).isTrue(); + assertThat(queue.getAvailableFuture().isDone()).isTrue(); + + putSuccessfully(queue, new StreamRecord<>(1, 0L)); + assertThat(queue.isAvailable()) + .as("queue with one free slot remaining should still be available") + .isTrue(); + + putSuccessfully(queue, new StreamRecord<>(2, 1L)); + assertThat(queue.isAvailable()).as("queue at capacity should be unavailable").isFalse(); + assertThat(queue.getAvailableFuture().isDone()).isFalse(); + } + + /** + * Tests that the future returned by {@link OrderedStreamElementQueue#getAvailableFuture()} + * completes once a slot is freed by emitting a completed head element, and that the queue + * becomes available again. + */ + @Test + void testGetAvailableFutureCompletesWhenSlotFreed() { + final OrderedStreamElementQueue<Integer> queue = new OrderedStreamElementQueue<>(2); + + ResultFuture<Integer> entry1 = putSuccessfully(queue, new StreamRecord<>(1, 0L)); + putSuccessfully(queue, new StreamRecord<>(2, 1L)); + + CompletableFuture<?> availableFuture = queue.getAvailableFuture(); + assertThat(availableFuture.isDone()).isFalse(); + assertThat(queue.isAvailable()).isFalse(); + + // Completing the head and emitting it frees one slot, which must complete the future + // captured while the queue was full. + entry1.complete(Collections.singleton(10)); + assertThat(popCompleted(queue)).containsExactly(new StreamRecord<>(10, 0L)); + + assertThat(availableFuture.isDone()) + .as("availability future should be completed when capacity is freed") + .isTrue(); + assertThat(queue.isAvailable()).isTrue(); + assertThat(queue.getAvailableFuture().isDone()).isTrue(); + } } diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueueTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueueTest.java index 35b8292aef4..c358552d3b2 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueueTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/queue/UnorderedStreamElementQueueTest.java @@ -25,6 +25,8 @@ import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; import org.junit.jupiter.api.Test; import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.CompletableFuture; import static org.apache.flink.streaming.api.operators.async.queue.QueueUtil.popCompleted; import static org.apache.flink.streaming.api.operators.async.queue.QueueUtil.putSuccessfully; @@ -93,4 +95,54 @@ class UnorderedStreamElementQueueTest { assertThat(queue.isEmpty()).isTrue(); assertThat(popCompleted(queue)).isEmpty(); } + + /** + * Tests that {@link UnorderedStreamElementQueue#isAvailable()} starts true, stays true while + * the queue has free capacity, and flips to false once it becomes full; the future returned by + * {@link UnorderedStreamElementQueue#getAvailableFuture()} mirrors this state. + */ + @Test + void testIsAvailableTogglesWithCapacity() { + final UnorderedStreamElementQueue<Integer> queue = new UnorderedStreamElementQueue<>(2); + + assertThat(queue.isAvailable()).isTrue(); + assertThat(queue.getAvailableFuture().isDone()).isTrue(); + + putSuccessfully(queue, new StreamRecord<>(1, 0L)); + assertThat(queue.isAvailable()) + .as("queue with one free slot remaining should still be available") + .isTrue(); + + putSuccessfully(queue, new StreamRecord<>(2, 1L)); + assertThat(queue.isAvailable()).as("queue at capacity should be unavailable").isFalse(); + assertThat(queue.getAvailableFuture().isDone()).isFalse(); + } + + /** + * Tests that the future returned by {@link UnorderedStreamElementQueue#getAvailableFuture()} + * completes once a slot is freed by emitting a completed element from the head segment, and + * that the queue becomes available again. + */ + @Test + void testGetAvailableFutureCompletesWhenSlotFreed() { + final UnorderedStreamElementQueue<Integer> queue = new UnorderedStreamElementQueue<>(2); + + ResultFuture<Integer> entry1 = putSuccessfully(queue, new StreamRecord<>(1, 0L)); + putSuccessfully(queue, new StreamRecord<>(2, 1L)); + + CompletableFuture<?> availableFuture = queue.getAvailableFuture(); + assertThat(availableFuture.isDone()).isFalse(); + assertThat(queue.isAvailable()).isFalse(); + + // Completing an entry in the head segment and emitting it frees one slot, which must + // complete the future captured while the queue was full. + entry1.complete(Collections.singleton(11)); + assertThat(popCompleted(queue)).containsExactly(new StreamRecord<>(11, 0L)); + + assertThat(availableFuture.isDone()) + .as("availability future should be completed when capacity is freed") + .isTrue(); + assertThat(queue.isAvailable()).isTrue(); + assertThat(queue.getAvailableFuture().isDone()).isTrue(); + } } diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskMailboxTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskMailboxTestHarness.java index be1934336c9..9dd7c85227c 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskMailboxTestHarness.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskMailboxTestHarness.java @@ -69,6 +69,10 @@ public class StreamTaskMailboxTestHarness<OUT> implements AutoCloseable { return streamTask; } + public OperatorChain<OUT, ?> getOperatorChain() { + return streamTask.operatorChain; + } + public TimerService getTimerService() { return streamTask.getTimerService(); }
