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 ba19adb4ff40ba02346a1a04d0add20fbb25253d Author: luogen.lg <[email protected]> AuthorDate: Wed Jun 3 16:09:30 2026 +0800 [FLINK-39835][runtime] Add SupportsChainAvailability and chain-level availability plumbing --- .../runtime/io/CompositeAvailabilityProvider.java | 51 ++++++++ .../api/operators/SupportsChainAvailability.java | 39 ++++++ .../streaming/runtime/io/RecordWriterOutput.java | 5 + .../streaming/runtime/tasks/OperatorChain.java | 66 +++++++++- .../flink/streaming/runtime/tasks/StreamTask.java | 22 +++- .../operators/StreamOperatorChainingTest.java | 142 +++++++++++++++++++++ .../tasks/CompositeAvailabilityProviderTest.java | 91 +++++++++++++ 7 files changed, 409 insertions(+), 7 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/CompositeAvailabilityProvider.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/CompositeAvailabilityProvider.java new file mode 100644 index 00000000000..1f6985e37be --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/CompositeAvailabilityProvider.java @@ -0,0 +1,51 @@ +/* + * 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.runtime.io; + +import org.apache.flink.annotation.Internal; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +/** Reports available only when all child {@link AvailabilityProvider}s are available. */ +@Internal +public class CompositeAvailabilityProvider implements AvailabilityProvider { + private final List<AvailabilityProvider> providers; + + private CompositeAvailabilityProvider(List<AvailabilityProvider> providers) { + this.providers = new ArrayList<>(providers); + } + + @Override + public CompletableFuture<?> getAvailableFuture() { + CompletableFuture<?> result = AvailabilityProvider.AVAILABLE; + for (AvailabilityProvider p : providers) { + result = AvailabilityProvider.and(result, p.getAvailableFuture()); + } + return result; + } + + public static AvailabilityProvider of(List<AvailabilityProvider> providers) { + if (providers.size() == 1) { + return providers.get(0); + } + return new CompositeAvailabilityProvider(providers); + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/SupportsChainAvailability.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/SupportsChainAvailability.java new file mode 100644 index 00000000000..269fc719190 --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/SupportsChainAvailability.java @@ -0,0 +1,39 @@ +/* + * 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.streaming.api.operators; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.runtime.io.AvailabilityProvider; + +/** + * Interface for operators that represent the availability of their chain segment. Such operators + * absorb downstream availability internally, handling back-pressure by deferring emits and yielding + * the mailbox rather than blocking. + * + * <p>{@code OperatorChain} injects the composite downstream {@link AvailabilityProvider} via {@link + * #setDownstreamAvailabilityProvider} and lets the operator represent the chain's availability. + */ +@Internal +public interface SupportsChainAvailability extends AvailabilityProvider { + + /** + * Called once by {@code OperatorChain} during construction to inject downstream availability. + */ + void setDownstreamAvailabilityProvider(AvailabilityProvider provider); +} diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/RecordWriterOutput.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/RecordWriterOutput.java index 815fa13bc39..3fe4aa1c108 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/RecordWriterOutput.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/RecordWriterOutput.java @@ -25,6 +25,7 @@ import org.apache.flink.metrics.SimpleCounter; import org.apache.flink.runtime.checkpoint.CheckpointException; import org.apache.flink.runtime.event.AbstractEvent; import org.apache.flink.runtime.event.WatermarkEvent; +import org.apache.flink.runtime.io.AvailabilityProvider; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; import org.apache.flink.runtime.io.network.api.writer.RecordWriter; import org.apache.flink.runtime.plugable.SerializationDelegate; @@ -222,6 +223,10 @@ public class RecordWriterOutput<OUT> recordWriter.flushAll(); } + public AvailabilityProvider getOutputAvailabilityProvider() { + return recordWriter; + } + @Override public void close() { recordWriter.close(); diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java index 78d202c62c7..6a3f0273cea 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java @@ -30,6 +30,8 @@ import org.apache.flink.runtime.checkpoint.CheckpointOptions; import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.event.AbstractEvent; import org.apache.flink.runtime.execution.Environment; +import org.apache.flink.runtime.io.AvailabilityProvider; +import org.apache.flink.runtime.io.CompositeAvailabilityProvider; import org.apache.flink.runtime.io.network.api.StopMode; import org.apache.flink.runtime.io.network.api.writer.RecordWriter; import org.apache.flink.runtime.io.network.api.writer.RecordWriterDelegate; @@ -61,6 +63,7 @@ import org.apache.flink.streaming.api.operators.StreamOperator; import org.apache.flink.streaming.api.operators.StreamOperatorFactory; import org.apache.flink.streaming.api.operators.StreamOperatorFactoryUtil; import org.apache.flink.streaming.api.operators.StreamTaskStateInitializer; +import org.apache.flink.streaming.api.operators.SupportsChainAvailability; import org.apache.flink.streaming.runtime.io.RecordWriterOutput; import org.apache.flink.streaming.runtime.io.StreamTaskSourceInput; import org.apache.flink.streaming.runtime.operators.sink.SinkWriterOperatorFactory; @@ -153,6 +156,8 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> protected boolean isClosed; + private AvailabilityProvider chainAvailabilityProvider = null; + public OperatorChain( StreamTask<OUT, OP> containingTask, RecordWriterDelegate<SerializationDelegate<StreamRecord<OUT>>> recordWriterDelegate) { @@ -198,6 +203,7 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> // we create the chain of operators and grab the collector that leads into the chain List<StreamOperatorWrapper<?, ?>> allOpWrappers = new ArrayList<>(chainedConfigs.size()); + List<AvailabilityProvider> chainAvailProviders = new ArrayList<>(); this.mainOperatorOutput = createOutputCollector( containingTask, @@ -207,7 +213,8 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> recordWriterOutputs, allOpWrappers, containingTask.getMailboxExecutorFactory(), - operatorFactory != null); + operatorFactory != null, + chainAvailProviders); if (operatorFactory != null) { Tuple2<OP, Optional<ProcessingTimeService>> mainOperatorAndTimeService = @@ -232,6 +239,26 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> mainOperatorAndTimeService.f1, true); + // Chain-availability head absorbs downstream availability. + if (mainOperator instanceof SupportsChainAvailability) { + SupportsChainAvailability head = (SupportsChainAvailability) mainOperator; + if (!chainAvailProviders.isEmpty()) { + head.setDownstreamAvailabilityProvider( + CompositeAvailabilityProvider.of(chainAvailProviders)); + } + chainAvailProviders = new ArrayList<>(); + chainAvailProviders.add(head); + } + + boolean hasChainAvailabilityOperator = + chainAvailProviders.stream() + .anyMatch( + provider -> provider instanceof SupportsChainAvailability); + if (hasChainAvailabilityOperator) { + chainAvailabilityProvider = + CompositeAvailabilityProvider.of(chainAvailProviders); + } + // add main operator to end of chain allOpWrappers.add(mainOperatorWrapper); @@ -485,6 +512,10 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> return isClosed; } + public AvailabilityProvider getChainAvailabilityProvider() { + return chainAvailabilityProvider; + } + /** Wrapper class to access the chained sources and their's outputs. */ public static class ChainedSource { private final WatermarkGaugeExposingOutput<StreamRecord<?>> chainedSourceOutput; @@ -712,7 +743,8 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> Map<IntermediateDataSetID, RecordWriterOutput<?>> recordWriterOutputs, List<StreamOperatorWrapper<?, ?>> allOperatorWrappers, MailboxExecutorFactory mailboxExecutorFactory, - boolean shouldAddMetric) { + boolean shouldAddMetric, + List<AvailabilityProvider> downstreamAvailProviders) { List<OutputWithChainingCheck<StreamRecord<T>>> allOutputs = new ArrayList<>(4); // create collectors for the network outputs @@ -723,6 +755,9 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> (RecordWriterOutput<T>) recordWriterOutputs.get(streamOutput.getDataSetId()); allOutputs.add(recordWriterOutput); + if (downstreamAvailProviders != null) { + downstreamAvailProviders.add(recordWriterOutput.getOutputAvailabilityProvider()); + } } // Create collectors for the chained outputs @@ -741,7 +776,8 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> allOperatorWrappers, outputEdge.getOutputTag(), mailboxExecutorFactory, - shouldAddMetric); + shouldAddMetric, + downstreamAvailProviders); checkState(output instanceof OutputWithChainingCheck); allOutputs.add((OutputWithChainingCheck) output); // If the operator has multiple downstream chained operators, only one of them should @@ -820,7 +856,10 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> List<StreamOperatorWrapper<?, ?>> allOperatorWrappers, OutputTag<IN> outputTag, MailboxExecutorFactory mailboxExecutorFactory, - boolean shouldAddMetricForPrevOperator) { + boolean shouldAddMetricForPrevOperator, + @Nullable List<AvailabilityProvider> parentDownstreamAvailProviders) { + List<AvailabilityProvider> myDownstreamAvailProviders = new ArrayList<>(); + // create the output that the operator writes to first. this may recursively create more // operators WatermarkGaugeExposingOutput<StreamRecord<OUT>> chainedOperatorOutput = @@ -832,7 +871,8 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> recordWriterOutputs, allOperatorWrappers, mailboxExecutorFactory, - true); + true, + myDownstreamAvailProviders); OneInputStreamOperator<IN, OUT> chainedOperator = createOperator( @@ -843,6 +883,22 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>> allOperatorWrappers, false); + // Chain-availability operator absorbs downstream; only it is surfaced upward. + if (chainedOperator instanceof SupportsChainAvailability) { + SupportsChainAvailability op = (SupportsChainAvailability) chainedOperator; + if (!myDownstreamAvailProviders.isEmpty()) { + op.setDownstreamAvailabilityProvider( + CompositeAvailabilityProvider.of(myDownstreamAvailProviders)); + } + if (parentDownstreamAvailProviders != null) { + parentDownstreamAvailProviders.add(op); + } + } else { + if (parentDownstreamAvailProviders != null) { + parentDownstreamAvailProviders.addAll(myDownstreamAvailProviders); + } + } + return wrapOperatorIntoOutput( chainedOperator, containingTask, diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java index d9d6738caf8..8c9270f581b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java @@ -349,6 +349,8 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>> @Nullable private final AvailabilityProvider changelogWriterAvailabilityProvider; + @Nullable private AvailabilityProvider chainAvailabilityProvider = null; + private long initializeStateEndTs; // ------------------------------------------------------------------------ @@ -721,11 +723,21 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>> } TaskIOMetricGroup ioMetrics = getEnvironment().getMetricGroup().getIOMetricGroup(); - PeriodTimer timer; + PeriodTimer timer = null; CompletableFuture<?> resumeFuture; - if (!recordWriter.isAvailable()) { + if (chainAvailabilityProvider == null && !recordWriter.isAvailable()) { + // No chain-availability operator; downstream is the direct backpressure source. timer = new GaugePeriodTimer(ioMetrics.getSoftBackPressuredTimePerSecond()); resumeFuture = recordWriter.getAvailableFuture(); + } else if (chainAvailabilityProvider != null && !chainAvailabilityProvider.isAvailable()) { + if (!recordWriter.isAvailable()) { + // Queue full AND downstream unavailable → backpressure (root cause: downstream). + timer = new GaugePeriodTimer(ioMetrics.getSoftBackPressuredTimePerSecond()); + resumeFuture = recordWriter.getAvailableFuture(); + } else { + // Queue full but downstream available → internal processing; fall back to busy. + resumeFuture = chainAvailabilityProvider.getAvailableFuture(); + } } else if (!inputProcessor.isAvailable()) { timer = new GaugePeriodTimer(ioMetrics.getIdleTimeMsPerSecond()); resumeFuture = inputProcessor.getAvailableFuture(); @@ -847,6 +859,7 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>> ? new FinishedOperatorChain<>(this, recordWriter) : new RegularOperatorChain<>(this, recordWriter); mainOperator = operatorChain.getMainOperator(); + chainAvailabilityProvider = operatorChain.getChainAvailabilityProvider(); getEnvironment() .getTaskStateManager() @@ -1361,6 +1374,11 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>> } private boolean taskIsAvailable() { + if (chainAvailabilityProvider != null) { + return chainAvailabilityProvider.isAvailable() + && (changelogWriterAvailabilityProvider == null + || changelogWriterAvailabilityProvider.isAvailable()); + } return recordWriter.isAvailable() && (changelogWriterAvailabilityProvider == null || changelogWriterAvailabilityProvider.isAvailable()); diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java index d425fc61b00..f3fc92d083b 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java @@ -19,8 +19,11 @@ package org.apache.flink.streaming.runtime.operators; import org.apache.flink.api.common.ExecutionConfig; +import org.apache.flink.api.common.functions.MapFunction; import org.apache.flink.configuration.Configuration; import org.apache.flink.runtime.execution.Environment; +import org.apache.flink.runtime.io.AvailabilityProvider; +import org.apache.flink.runtime.io.network.api.writer.AvailabilityTestResultPartitionWriter; import org.apache.flink.runtime.io.network.api.writer.RecordWriterDelegate; import org.apache.flink.runtime.jobgraph.JobGraph; import org.apache.flink.runtime.jobgraph.JobVertex; @@ -37,6 +40,7 @@ import org.apache.flink.streaming.api.graph.StreamConfig; import org.apache.flink.streaming.api.operators.StreamMap; import org.apache.flink.streaming.api.operators.StreamOperator; import org.apache.flink.streaming.api.operators.StreamTaskStateInitializer; +import org.apache.flink.streaming.api.operators.SupportsChainAvailability; import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; import org.apache.flink.streaming.runtime.tasks.OperatorChain; import org.apache.flink.streaming.runtime.tasks.RegularOperatorChain; @@ -49,7 +53,9 @@ import org.apache.flink.util.OutputTag; import org.junit.jupiter.api.Test; import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.concurrent.CompletableFuture; import static org.apache.flink.streaming.api.operators.StreamOperatorUtils.setupStreamOperator; import static org.assertj.core.api.Assertions.assertThat; @@ -276,6 +282,142 @@ class StreamOperatorChainingTest { } } + /** + * Exercises the head branch of {@link OperatorChain}: when the main operator implements {@link + * SupportsChainAvailability}, the chain availability provider must be the head operator itself. + */ + @Test + void testChainAvailabilityProviderForHeadChainAvailabilityOperator() throws Exception { + JobVertex chainedVertex = buildChainedVertexWithChainAvailabilityOperator(true); + + StreamConfig streamConfig = new StreamConfig(chainedVertex.getConfiguration()); + + try (MockEnvironment environment = createMockEnvironment(chainedVertex.getName())) { + environment.addOutputs( + Collections.singletonList(new AvailabilityTestResultPartitionWriter(true))); + StreamTask<Integer, ChainAvailabilityMapOperator<Integer, Integer>> mockTask = + createMockTask(streamConfig, environment); + OperatorChain<Integer, ChainAvailabilityMapOperator<Integer, Integer>> operatorChain = + createOperatorChain(streamConfig, environment, mockTask); + + AvailabilityProvider chainAvailability = operatorChain.getChainAvailabilityProvider(); + assertThat(chainAvailability) + .as("Head chain-availability operator should represent chain" + " availability") + .isSameAs(operatorChain.getMainOperator()); + assertThat(chainAvailability.isAvailable()).as("Operator starts available").isTrue(); + } + } + + /** + * Exercises the chained branch of {@link OperatorChain}: when a {@link + * SupportsChainAvailability} operator sits behind the head, it replaces the downstream leaves + * in the parent's collector, so the chain availability provider must be that chained operator. + */ + @Test + void testChainAvailabilityProviderForChainedChainAvailabilityOperator() throws Exception { + JobVertex chainedVertex = buildChainedVertexWithChainAvailabilityOperator(false); + + StreamConfig streamConfig = new StreamConfig(chainedVertex.getConfiguration()); + + try (MockEnvironment environment = createMockEnvironment(chainedVertex.getName())) { + environment.addOutputs( + Collections.singletonList(new AvailabilityTestResultPartitionWriter(true))); + StreamTask<Integer, StreamMap<Integer, Integer>> mockTask = + createMockTask(streamConfig, environment); + OperatorChain<Integer, StreamMap<Integer, Integer>> operatorChain = + createOperatorChain(streamConfig, environment, mockTask); + + SupportsChainAvailability chainedOp = + findSingleChainAvailabilityOperator(operatorChain); + AvailabilityProvider chainAvailability = operatorChain.getChainAvailabilityProvider(); + assertThat(chainAvailability) + .as( + "Chained chain-availability operator should replace downstream" + + " leaves and act as chain availability boundary") + .isSameAs(chainedOp); + assertThat(chainAvailability) + .as( + "Head operator (not the chain-availability one) is not the chain" + + " availability provider") + .isNotSameAs(operatorChain.getMainOperator()); + assertThat(chainAvailability.isAvailable()).as("Operator starts available").isTrue(); + } + } + + private static JobVertex buildChainedVertexWithChainAvailabilityOperator( + boolean chainAvailabilityAtHead) { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(2); + + DataStream<Integer> input = env.fromData(1, 2, 3); + + if (chainAvailabilityAtHead) { + input = + input.transform( + "ChainAvailabilityMap", + input.getType(), + new ChainAvailabilityMapOperator<>(value -> value)); + input = input.map(value -> value); + } else { + input = input.map(value -> value); + input = + input.transform( + "ChainAvailabilityMap", + input.getType(), + new ChainAvailabilityMapOperator<>(value -> value)); + } + + input.map(value -> value) + .startNewChain() + .addSink( + new SinkFunction<Integer>() { + @Override + public void invoke(Integer value, Context ctx) {} + }); + + JobGraph jobGraph = env.getStreamGraph().getJobGraph(); + assertThat(jobGraph.getVerticesSortedTopologicallyFromSources()).hasSize(3); + return jobGraph.getVerticesSortedTopologicallyFromSources().get(1); + } + + private static SupportsChainAvailability findSingleChainAvailabilityOperator( + OperatorChain<?, ?> chain) { + SupportsChainAvailability found = null; + for (StreamOperatorWrapper<?, ?> wrapper : chain.getAllOperators()) { + if (wrapper.getStreamOperator() instanceof SupportsChainAvailability) { + assertThat(found) + .as("Expected exactly one SupportsChainAvailability operator in chain") + .isNull(); + found = (SupportsChainAvailability) wrapper.getStreamOperator(); + } + } + assertThat(found).as("Expected a SupportsChainAvailability operator in chain").isNotNull(); + return found; + } + + /** + * A {@link StreamMap} that also implements {@link SupportsChainAvailability} for testing + * OperatorChain availability-provider wiring without depending on flink-streaming-java + * operators. + */ + private static class ChainAvailabilityMapOperator<IN, OUT> extends StreamMap<IN, OUT> + implements SupportsChainAvailability { + + private static final long serialVersionUID = 1L; + + ChainAvailabilityMapOperator(MapFunction<IN, OUT> mapper) { + super(mapper); + } + + @Override + public void setDownstreamAvailabilityProvider(AvailabilityProvider provider) {} + + @Override + public CompletableFuture<?> getAvailableFuture() { + return AvailabilityProvider.AVAILABLE; + } + } + private <IN, OT extends StreamOperator<IN>> OperatorChain<IN, OT> createOperatorChain( StreamConfig streamConfig, Environment environment, StreamTask<IN, OT> task) { return new TestOperatorChain<>( diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/CompositeAvailabilityProviderTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/CompositeAvailabilityProviderTest.java new file mode 100644 index 00000000000..6e4fd5e04ee --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/CompositeAvailabilityProviderTest.java @@ -0,0 +1,91 @@ +/* + * 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.streaming.runtime.tasks; + +import org.apache.flink.runtime.io.AvailabilityProvider; +import org.apache.flink.runtime.io.CompositeAvailabilityProvider; + +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.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link CompositeAvailabilityProvider}. */ +class CompositeAvailabilityProviderTest { + + @Test + void testSingleChildReturnsChildDirectly() { + TestAvailabilityProvider child = new TestAvailabilityProvider(); + AvailabilityProvider composite = + CompositeAvailabilityProvider.of(Collections.singletonList(child)); + + assertThat(composite).isSameAs(child); + assertThat(composite.getAvailableFuture()).isSameAs(child.future); + assertThat(composite.isAvailable()).isFalse(); + + child.future.complete(null); + assertThat(composite.isAvailable()).isTrue(); + } + + @Test + void testMultipleChildrenWrappedInComposite() { + TestAvailabilityProvider first = new TestAvailabilityProvider(); + TestAvailabilityProvider second = new TestAvailabilityProvider(); + List<AvailabilityProvider> children = Arrays.asList(first, second); + + AvailabilityProvider composite = CompositeAvailabilityProvider.of(children); + + assertThat(composite).isInstanceOf(CompositeAvailabilityProvider.class); + assertThat(composite.isAvailable()).isFalse(); + } + + @Test + void testGetAvailableFutureCompletesOnlyWhenAllChildrenComplete() { + TestAvailabilityProvider first = new TestAvailabilityProvider(); + TestAvailabilityProvider second = new TestAvailabilityProvider(); + AvailabilityProvider composite = + CompositeAvailabilityProvider.of(Arrays.asList(first, second)); + + CompletableFuture<?> available = composite.getAvailableFuture(); + assertThat(available.isDone()).isFalse(); + + first.future.complete(null); + assertThat(available.isDone()) + .as("composite future should remain incomplete while one child is pending") + .isFalse(); + + second.future.complete(null); + assertThat(available.isDone()).isTrue(); + assertThat(composite.isAvailable()).isTrue(); + } + + /** Minimal {@link AvailabilityProvider} backed by a mutable {@link CompletableFuture}. */ + private static final class TestAvailabilityProvider implements AvailabilityProvider { + private final CompletableFuture<Void> future = new CompletableFuture<>(); + + @Override + public CompletableFuture<?> getAvailableFuture() { + return future; + } + } +}
