[FLINK-13249][runtime] Fix handling of partition producer responses b… (apache#9138)

* [FLINK-13249][runtime] Fix handling of partition producer responses by running them with the task's executor

* Review comments

Author: StefanRRichter

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionProducerStateProvider.java

@@ -22,7 +22,7 @@
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.types.Either;
 
-import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 
 /**
  * Request execution state of partition producer, the response accepts state check callbacks.
@@ -34,11 +34,12 @@ public interface PartitionProducerStateProvider {
 	 * @param intermediateDataSetId ID of the parent intermediate data set.
 	 * @param resultPartitionId ID of the result partition to check. This
 	 * identifies the producing execution and partition.
-	 * @return a future with response handle.
+	 * @param responseConsumer consumer for the response handle.
 	 */
-	CompletableFuture<? extends ResponseHandle> requestPartitionProducerState(
+	void requestPartitionProducerState(
 		IntermediateDataSetID intermediateDataSetId,
-		ResultPartitionID resultPartitionId);
+		ResultPartitionID resultPartitionId,
+		Consumer<? super ResponseHandle> responseConsumer);
 
 	/**
 	 * Result of state query, accepts state check callbacks.

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGate.java

@@ -601,8 +601,8 @@ void notifyChannelNonEmpty(InputChannel channel) {
 	void triggerPartitionStateCheck(ResultPartitionID partitionId) {
 		partitionProducerStateProvider.requestPartitionProducerState(
 			consumedResultId,
-			partitionId)
-			.thenAccept(responseHandle -> {
+			partitionId,
+			((PartitionProducerStateProvider.ResponseHandle responseHandle) -> {
 				boolean isProducingState = new RemoteChannelStateChecker(partitionId, owningTaskName)
 					.isProducerReadyOrAbortConsumption(responseHandle);
 				if (isProducingState) {
@@ -612,7 +612,7 @@ void triggerPartitionStateCheck(ResultPartitionID partitionId) {
 						responseHandle.failConsumption(t);
 					}
 				}
-			});
+			}));
 	}
 
 	private void queueChannel(InputChannel channel) {

flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java

@@ -99,6 +99,7 @@
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
+import java.util.function.Consumer;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -1080,18 +1081,21 @@ else if (current == ExecutionState.RUNNING) {
 	// ------------------------------------------------------------------------
 
 	@Override
-	public CompletableFuture<PartitionProducerStateResponseHandle> requestPartitionProducerState(
+	public void requestPartitionProducerState(
 			final IntermediateDataSetID intermediateDataSetId,
-			final ResultPartitionID resultPartitionId) {
+			final ResultPartitionID resultPartitionId,
+			Consumer<? super ResponseHandle> responseConsumer) {
+
 		final CompletableFuture<ExecutionState> futurePartitionState =
 			partitionProducerStateChecker.requestPartitionProducerState(
 				jobId,
 				intermediateDataSetId,
 				resultPartitionId);
-		final CompletableFuture<PartitionProducerStateResponseHandle> result =
-			futurePartitionState.handleAsync(PartitionProducerStateResponseHandle::new, executor);
-		FutureUtils.assertNoException(result);
-		return result;
+
+		FutureUtils.assertNoException(
+			futurePartitionState
+				.handle(PartitionProducerStateResponseHandle::new)
+				.thenAcceptAsync(responseConsumer, executor));
 	}
 
 	// ------------------------------------------------------------------------

flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateBuilder.java

@@ -21,23 +21,19 @@
 import org.apache.flink.runtime.io.network.NettyShuffleEnvironment;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
 import org.apache.flink.runtime.io.network.partition.PartitionProducerStateProvider;
-import org.apache.flink.runtime.io.network.partition.PartitionProducerStateProvider.ResponseHandle;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.runtime.taskmanager.NettyShuffleEnvironmentConfiguration;
 import org.apache.flink.util.function.SupplierWithException;
 
 import java.io.IOException;
-import java.util.concurrent.CompletableFuture;
 
 /**
  * Utility class to encapsulate the logic of building a {@link SingleInputGate} instance.
  */
 public class SingleInputGateBuilder {
 
-	private static final CompletableFuture<ResponseHandle> NO_OP_PRODUCER_CHECKER_RESULT = new CompletableFuture<>();
-
-	public static final PartitionProducerStateProvider NO_OP_PRODUCER_CHECKER = (dsid, id) -> NO_OP_PRODUCER_CHECKER_RESULT;
+	public static final PartitionProducerStateProvider NO_OP_PRODUCER_CHECKER = (dsid, id, consumer) -> {};
 
 	private final IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
 

flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java

@@ -656,7 +656,7 @@ public void testTriggerPartitionStateUpdate() throws Exception {
 			final CompletableFuture<ExecutionState> promise = new CompletableFuture<>();
 			when(partitionChecker.requestPartitionProducerState(eq(task.getJobID()), eq(resultId), eq(partitionId))).thenReturn(promise);
 
-			task.requestPartitionProducerState(resultId, partitionId).thenAccept(checkResult ->
+			task.requestPartitionProducerState(resultId, partitionId, checkResult ->
 				assertThat(remoteChannelStateChecker.isProducerReadyOrAbortConsumption(checkResult), is(false))
 			);
 
@@ -680,7 +680,7 @@ public void testTriggerPartitionStateUpdate() throws Exception {
 			final CompletableFuture<ExecutionState> promise = new CompletableFuture<>();
 			when(partitionChecker.requestPartitionProducerState(eq(task.getJobID()), eq(resultId), eq(partitionId))).thenReturn(promise);
 
-			task.requestPartitionProducerState(resultId, partitionId).thenAccept(checkResult ->
+			task.requestPartitionProducerState(resultId, partitionId, checkResult ->
 				assertThat(remoteChannelStateChecker.isProducerReadyOrAbortConsumption(checkResult), is(false))
 			);
 
@@ -711,7 +711,7 @@ public void testTriggerPartitionStateUpdate() throws Exception {
 				CompletableFuture<ExecutionState> promise = new CompletableFuture<>();
 				when(partitionChecker.requestPartitionProducerState(eq(task.getJobID()), eq(resultId), eq(partitionId))).thenReturn(promise);
 
-				task.requestPartitionProducerState(resultId, partitionId).thenAccept(checkResult -> {
+				task.requestPartitionProducerState(resultId, partitionId, checkResult -> {
 					if (remoteChannelStateChecker.isProducerReadyOrAbortConsumption(checkResult)) {
 						callCount.incrementAndGet();
 					}
@@ -749,7 +749,7 @@ public void testTriggerPartitionStateUpdate() throws Exception {
 				CompletableFuture<ExecutionState> promise = new CompletableFuture<>();
 				when(partitionChecker.requestPartitionProducerState(eq(task.getJobID()), eq(resultId), eq(partitionId))).thenReturn(promise);
 
-				task.requestPartitionProducerState(resultId, partitionId).thenAccept(checkResult -> {
+				task.requestPartitionProducerState(resultId, partitionId, checkResult -> {
 					if (remoteChannelStateChecker.isProducerReadyOrAbortConsumption(checkResult)) {
 						callCount.incrementAndGet();
 					}