You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ga...@apache.org on 2019/06/13 19:11:33 UTC
[flink] branch master updated (27a1c41 -> 8a39c84)
This is an automated email from the ASF dual-hosted git repository.
gary pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git.
from 27a1c41 [FLINK-12822] Add explicit transformer from SerializableOptional to Optional
new 90f2066 [FLINK-12608][runtime] Add getVertexOrThrow and getResultPartitionOrThrow to SchedulingTopology
new 8a39c84 [FLINK-12372] [runtime] implement ExecutionSlotAllocator
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../runtime/executiongraph/ExecutionVertex.java | 2 +-
.../scheduler/DefaultExecutionSlotAllocator.java | 242 +++++++++++++++
...mentOption.java => ExecutionSlotAllocator.java} | 45 +--
.../ExecutionVertexSchedulingRequirements.java | 160 ++++++++++
.../scheduler/InputsLocationsRetriever.java | 49 +++
...ion.java => SlotExecutionVertexAssignment.java} | 20 +-
.../LazyFromSourcesSchedulingStrategy.java | 17 +-
.../scheduler/strategy/SchedulingTopology.java | 24 ++
...ecutionSlotAllocatorPreferredLocationsTest.java | 116 +++++++
.../DefaultExecutionSlotAllocatorTest.java | 336 +++++++++++++++++++++
.../scheduler/TestingInputsLocationsRetriever.java | 94 ++++++
...cutionGraphToSchedulingTopologyAdapterTest.java | 22 ++
12 files changed, 1087 insertions(+), 40 deletions(-)
create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
copy flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/{ExecutionVertexDeploymentOption.java => ExecutionSlotAllocator.java} (53%)
create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionVertexSchedulingRequirements.java
create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/InputsLocationsRetriever.java
copy flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/{ExecutionVertexDeploymentOption.java => SlotExecutionVertexAssignment.java} (66%)
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorPreferredLocationsTest.java
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorTest.java
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingInputsLocationsRetriever.java
[flink] 01/02: [FLINK-12608][runtime] Add getVertexOrThrow and
getResultPartitionOrThrow to SchedulingTopology
Posted by ga...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
gary pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 90f2066c9d5ebf5da9749ae308086ca427d0e043
Author: Bo WANG <wb...@gmail.com>
AuthorDate: Tue Jun 4 09:59:18 2019 +0800
[FLINK-12608][runtime] Add getVertexOrThrow and getResultPartitionOrThrow to SchedulingTopology
This closes #8603.
---
.../LazyFromSourcesSchedulingStrategy.java | 17 +++++----------
.../scheduler/strategy/SchedulingTopology.java | 24 ++++++++++++++++++++++
...cutionGraphToSchedulingTopologyAdapterTest.java | 22 ++++++++++++++++++++
3 files changed, 51 insertions(+), 12 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/LazyFromSourcesSchedulingStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/LazyFromSourcesSchedulingStrategy.java
index b5a09d4..55a77c8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/LazyFromSourcesSchedulingStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/LazyFromSourcesSchedulingStrategy.java
@@ -86,7 +86,7 @@ public class LazyFromSourcesSchedulingStrategy implements SchedulingStrategy {
// increase counter of the dataset first
verticesToRestart
.stream()
- .map(this::getSchedulingVertex)
+ .map(schedulingTopology::getVertexOrThrow)
.flatMap(vertex -> vertex.getProducedResultPartitions().stream())
.forEach(inputConstraintChecker::resetSchedulingResultPartition);
@@ -99,7 +99,7 @@ public class LazyFromSourcesSchedulingStrategy implements SchedulingStrategy {
return;
}
- final Set<SchedulingExecutionVertex> verticesToSchedule = getSchedulingVertex(executionVertexId)
+ final Set<SchedulingExecutionVertex> verticesToSchedule = schedulingTopology.getVertexOrThrow(executionVertexId)
.getProducedResultPartitions()
.stream()
.flatMap(partition -> inputConstraintChecker.markSchedulingResultPartitionFinished(partition).stream())
@@ -112,15 +112,13 @@ public class LazyFromSourcesSchedulingStrategy implements SchedulingStrategy {
@Override
public void onPartitionConsumable(ExecutionVertexID executionVertexId, ResultPartitionID resultPartitionId) {
final SchedulingResultPartition resultPartition = schedulingTopology
- .getResultPartition(resultPartitionId.getPartitionId())
- .orElseThrow(() -> new IllegalStateException("can not find scheduling result partition for "
- + resultPartitionId));
+ .getResultPartitionOrThrow(resultPartitionId.getPartitionId());
if (!resultPartition.getPartitionType().isPipelined()) {
return;
}
- final SchedulingExecutionVertex producerVertex = getSchedulingVertex(executionVertexId);
+ final SchedulingExecutionVertex producerVertex = schedulingTopology.getVertexOrThrow(executionVertexId);
if (!producerVertex.getProducedResultPartitions().contains(resultPartition)) {
throw new IllegalStateException("partition " + resultPartitionId
+ " is not the produced partition of " + executionVertexId);
@@ -129,16 +127,11 @@ public class LazyFromSourcesSchedulingStrategy implements SchedulingStrategy {
allocateSlotsAndDeployExecutionVertices(resultPartition.getConsumers());
}
- private SchedulingExecutionVertex getSchedulingVertex(final ExecutionVertexID executionVertexId) {
- return schedulingTopology.getVertex(executionVertexId)
- .orElseThrow(() -> new IllegalStateException("can not find scheduling vertex for " + executionVertexId));
- }
-
private void allocateSlotsAndDeployExecutionVertexIds(Set<ExecutionVertexID> verticesToSchedule) {
allocateSlotsAndDeployExecutionVertices(
verticesToSchedule
.stream()
- .map(this::getSchedulingVertex)
+ .map(schedulingTopology::getVertexOrThrow)
.collect(Collectors.toList()));
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/SchedulingTopology.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/SchedulingTopology.java
index eb117d2..ce037a3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/SchedulingTopology.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/strategy/SchedulingTopology.java
@@ -44,10 +44,34 @@ public interface SchedulingTopology {
Optional<SchedulingExecutionVertex> getVertex(ExecutionVertexID executionVertexId);
/**
+ * Looks up the {@link SchedulingExecutionVertex} for the given {@link ExecutionVertexID}.
+ *
+ * @param executionVertexId identifying the respective scheduling vertex
+ * @return The respective scheduling vertex
+ * @throws IllegalArgumentException If the vertex does not exist
+ */
+ default SchedulingExecutionVertex getVertexOrThrow(ExecutionVertexID executionVertexId) {
+ return getVertex(executionVertexId).orElseThrow(
+ () -> new IllegalArgumentException("can not find vertex: " + executionVertexId));
+ }
+
+ /**
* Looks up the {@link SchedulingResultPartition} for the given {@link IntermediateResultPartitionID}.
*
* @param intermediateResultPartitionId identifying the respective scheduling result partition
* @return Optional containing the respective scheduling result partition or none if the partition does not exist
*/
Optional<SchedulingResultPartition> getResultPartition(IntermediateResultPartitionID intermediateResultPartitionId);
+
+ /**
+ * Looks up the {@link SchedulingResultPartition} for the given {@link IntermediateResultPartitionID}.
+ *
+ * @param intermediateResultPartitionId identifying the respective scheduling result partition
+ * @return The respective scheduling result partition
+ * @throws IllegalArgumentException If the partition does not exist
+ */
+ default SchedulingResultPartition getResultPartitionOrThrow(IntermediateResultPartitionID intermediateResultPartitionId) {
+ return getResultPartition(intermediateResultPartitionId).orElseThrow(
+ () -> new IllegalArgumentException("can not find partition: " + intermediateResultPartitionId));
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adapter/ExecutionGraphToSchedulingTopologyAdapterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adapter/ExecutionGraphToSchedulingTopologyAdapterTest.java
index b20a2f4..5a925da 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adapter/ExecutionGraphToSchedulingTopologyAdapterTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adapter/ExecutionGraphToSchedulingTopologyAdapterTest.java
@@ -27,6 +27,7 @@ import org.apache.flink.runtime.executiongraph.TestRestartStrategy;
import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
import org.apache.flink.runtime.jobgraph.JobVertex;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.scheduler.strategy.SchedulingExecutionVertex;
import org.apache.flink.runtime.scheduler.strategy.SchedulingResultPartition;
@@ -53,6 +54,7 @@ import static org.apache.flink.runtime.io.network.partition.ResultPartitionType.
import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.fail;
/**
* Unit tests for {@link ExecutionGraphToSchedulingTopologyAdapter}.
@@ -103,6 +105,26 @@ public class ExecutionGraphToSchedulingTopologyAdapterTest extends TestLogger {
}
}
+ @Test
+ public void testGetVertexOrThrow() {
+ try {
+ adapter.getVertexOrThrow(new ExecutionVertexID(new JobVertexID(), 0));
+ fail("get not exist vertex");
+ } catch (IllegalArgumentException exception) {
+ // expected
+ }
+ }
+
+ @Test
+ public void testResultPartitionOrThrow() {
+ try {
+ adapter.getResultPartitionOrThrow(new IntermediateResultPartitionID());
+ fail("get not exist result partition");
+ } catch (IllegalArgumentException exception) {
+ // expected
+ }
+ }
+
private static void assertGraphEquals(
ExecutionGraph originalGraph,
SchedulingTopology adaptedTopology) {
[flink] 02/02: [FLINK-12372] [runtime] implement
ExecutionSlotAllocator
Posted by ga...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
gary pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8a39c84e29c9ad64febbd850c99b40d28b5359d1
Author: shuai-xu <sh...@foxmail.com>
AuthorDate: Mon May 20 17:15:41 2019 +0800
[FLINK-12372] [runtime] implement ExecutionSlotAllocator
This closes #8486.
---
.../runtime/executiongraph/ExecutionVertex.java | 2 +-
.../scheduler/DefaultExecutionSlotAllocator.java | 242 +++++++++++++++
.../runtime/scheduler/ExecutionSlotAllocator.java | 51 ++++
.../ExecutionVertexSchedulingRequirements.java | 160 ++++++++++
.../scheduler/InputsLocationsRetriever.java | 49 +++
.../scheduler/SlotExecutionVertexAssignment.java | 52 ++++
...ecutionSlotAllocatorPreferredLocationsTest.java | 116 +++++++
.../DefaultExecutionSlotAllocatorTest.java | 336 +++++++++++++++++++++
.../scheduler/TestingInputsLocationsRetriever.java | 94 ++++++
9 files changed, 1101 insertions(+), 1 deletion(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index 9b0cfa0..0d01e25 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -71,7 +71,7 @@ public class ExecutionVertex implements AccessExecutionVertex, Archiveable<Archi
private static final Logger LOG = ExecutionGraph.LOG;
- private static final int MAX_DISTINCT_LOCATIONS_TO_CONSIDER = 8;
+ public static final int MAX_DISTINCT_LOCATIONS_TO_CONSIDER = 8;
// --------------------------------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
new file mode 100644
index 0000000..7948408
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
@@ -0,0 +1,242 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.time.Time;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.SlotProfile;
+import org.apache.flink.runtime.concurrent.FutureUtils;
+import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
+import org.apache.flink.runtime.jobmanager.scheduler.ScheduledUnit;
+import org.apache.flink.runtime.jobmaster.LogicalSlot;
+import org.apache.flink.runtime.jobmaster.SlotRequestId;
+import org.apache.flink.runtime.jobmaster.slotpool.SlotProvider;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+import static org.apache.flink.runtime.executiongraph.ExecutionVertex.MAX_DISTINCT_LOCATIONS_TO_CONSIDER;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * Default {@link ExecutionSlotAllocator} which will use {@link SlotProvider} to allocate slots and
+ * keep the unfulfilled requests for further cancellation.
+ */
+public class DefaultExecutionSlotAllocator implements ExecutionSlotAllocator {
+
+ private static final Logger LOG = LoggerFactory.getLogger(DefaultExecutionSlotAllocator.class);
+
+ /**
+ * Store the uncompleted slot assignments.
+ */
+ private final Map<ExecutionVertexID, SlotExecutionVertexAssignment> pendingSlotAssignments;
+
+ private final SlotProvider slotProvider;
+
+ private final InputsLocationsRetriever inputsLocationsRetriever;
+
+ private final Time allocationTimeout;
+
+ public DefaultExecutionSlotAllocator(
+ SlotProvider slotProvider,
+ InputsLocationsRetriever inputsLocationsRetriever,
+ Time allocationTimeout) {
+ this.slotProvider = checkNotNull(slotProvider);
+ this.inputsLocationsRetriever = checkNotNull(inputsLocationsRetriever);
+ this.allocationTimeout = checkNotNull(allocationTimeout);
+
+ pendingSlotAssignments = new HashMap<>();
+ }
+
+ @Override
+ public Collection<SlotExecutionVertexAssignment> allocateSlotsFor(
+ Collection<ExecutionVertexSchedulingRequirements> executionVertexSchedulingRequirements) {
+
+ List<SlotExecutionVertexAssignment> slotExecutionVertexAssignments =
+ new ArrayList<>(executionVertexSchedulingRequirements.size());
+
+ Set<AllocationID> allPreviousAllocationIds = computeAllPriorAllocationIds(executionVertexSchedulingRequirements);
+
+ for (ExecutionVertexSchedulingRequirements schedulingRequirements : executionVertexSchedulingRequirements) {
+ final ExecutionVertexID executionVertexId = schedulingRequirements.getExecutionVertexId();
+ final SlotRequestId slotRequestId = new SlotRequestId();
+ final SlotSharingGroupId slotSharingGroupId = schedulingRequirements.getSlotSharingGroupId();
+
+ LOG.debug("Allocate slot with id {} for execution {}", slotRequestId, executionVertexId);
+
+ CompletableFuture<LogicalSlot> slotFuture = calculatePreferredLocations(
+ executionVertexId,
+ schedulingRequirements.getPreferredLocations(),
+ inputsLocationsRetriever).thenCompose(
+ (Collection<TaskManagerLocation> preferredLocations) ->
+ slotProvider.allocateSlot(
+ slotRequestId,
+ new ScheduledUnit(
+ executionVertexId.getJobVertexId(),
+ slotSharingGroupId,
+ schedulingRequirements.getCoLocationConstraint()),
+ new SlotProfile(
+ schedulingRequirements.getResourceProfile(),
+ preferredLocations,
+ Arrays.asList(schedulingRequirements.getPreviousAllocationId()),
+ allPreviousAllocationIds),
+ true,
+ allocationTimeout));
+
+ SlotExecutionVertexAssignment slotExecutionVertexAssignment =
+ new SlotExecutionVertexAssignment(executionVertexId, slotFuture);
+ // add to map first to avoid the future completed before added.
+ pendingSlotAssignments.put(executionVertexId, slotExecutionVertexAssignment);
+
+ slotFuture.whenComplete(
+ (ignored, throwable) -> {
+ pendingSlotAssignments.remove(executionVertexId);
+ if (throwable != null) {
+ slotProvider.cancelSlotRequest(slotRequestId, slotSharingGroupId, throwable);
+ }
+ });
+
+ slotExecutionVertexAssignments.add(slotExecutionVertexAssignment);
+ }
+
+ return slotExecutionVertexAssignments;
+ }
+
+ @Override
+ public void cancel(ExecutionVertexID executionVertexId) {
+ SlotExecutionVertexAssignment slotExecutionVertexAssignment = pendingSlotAssignments.get(executionVertexId);
+ if (slotExecutionVertexAssignment != null) {
+ slotExecutionVertexAssignment.getLogicalSlotFuture().cancel(false);
+ }
+ }
+
+ @Override
+ public CompletableFuture<Void> stop() {
+ List<ExecutionVertexID> executionVertexIds = new ArrayList<>(pendingSlotAssignments.keySet());
+ executionVertexIds.forEach(this::cancel);
+
+ return CompletableFuture.completedFuture(null);
+ }
+
+ /**
+ * Calculates the preferred locations for an execution.
+ * It will first try to use preferred locations based on state,
+ * if null, will use the preferred locations based on inputs.
+ */
+ private static CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(
+ ExecutionVertexID executionVertexId,
+ Collection<TaskManagerLocation> preferredLocationsBasedOnState,
+ InputsLocationsRetriever inputsLocationsRetriever) {
+
+ if (!preferredLocationsBasedOnState.isEmpty()) {
+ return CompletableFuture.completedFuture(preferredLocationsBasedOnState);
+ }
+
+ return getPreferredLocationsBasedOnInputs(executionVertexId, inputsLocationsRetriever);
+ }
+
+ /**
+ * Gets the location preferences of the execution, as determined by the locations
+ * of the predecessors from which it receives input data.
+ * If there are more than {@link ExecutionVertex#MAX_DISTINCT_LOCATIONS_TO_CONSIDER} different locations of source data,
+ * or neither the sources have not been started nor will be started with the execution together,
+ * this method returns an empty collection to indicate no location preference.
+ *
+ * @return The preferred locations based in input streams, or an empty iterable,
+ * if there is no input-based preference.
+ */
+ @VisibleForTesting
+ static CompletableFuture<Collection<TaskManagerLocation>> getPreferredLocationsBasedOnInputs(
+ ExecutionVertexID executionVertexId,
+ InputsLocationsRetriever inputsLocationsRetriever) {
+ CompletableFuture<Collection<TaskManagerLocation>> preferredLocations =
+ CompletableFuture.completedFuture(Collections.emptyList());
+
+ Collection<CompletableFuture<TaskManagerLocation>> locationsFutures = new ArrayList<>();
+
+ Collection<Collection<ExecutionVertexID>> allProducers =
+ inputsLocationsRetriever.getConsumedResultPartitionsProducers(executionVertexId);
+ for (Collection<ExecutionVertexID> producers : allProducers) {
+
+ for (ExecutionVertexID producer : producers) {
+ Optional<CompletableFuture<TaskManagerLocation>> optionalLocationFuture =
+ inputsLocationsRetriever.getTaskManagerLocation(producer);
+ optionalLocationFuture.ifPresent(locationsFutures::add);
+ // If the parallelism is large, wait for all futures coming back may cost a long time.
+ if (locationsFutures.size() > MAX_DISTINCT_LOCATIONS_TO_CONSIDER) {
+ locationsFutures.clear();
+ break;
+ }
+ }
+
+ CompletableFuture<Collection<TaskManagerLocation>> uniqueLocationsFuture =
+ FutureUtils.combineAll(locationsFutures).thenApply(HashSet::new);
+ preferredLocations = preferredLocations.thenCombine(
+ uniqueLocationsFuture,
+ (locationsOnOneEdge, locationsOnAnotherEdge) -> {
+ if ((!locationsOnOneEdge.isEmpty() && locationsOnAnotherEdge.size() > locationsOnOneEdge.size())
+ || locationsOnAnotherEdge.isEmpty()) {
+ return locationsOnOneEdge;
+ } else {
+ return locationsOnAnotherEdge;
+ }
+ });
+ locationsFutures.clear();
+ }
+ return preferredLocations;
+ }
+
+ /**
+ * Computes and returns a set with the prior allocation ids from all execution vertices scheduled together.
+ *
+ * @param executionVertexSchedulingRequirements contains the execution vertices which are scheduled together
+ */
+ @VisibleForTesting
+ static Set<AllocationID> computeAllPriorAllocationIds(
+ Collection<ExecutionVertexSchedulingRequirements> executionVertexSchedulingRequirements) {
+ return executionVertexSchedulingRequirements
+ .stream()
+ .map(ExecutionVertexSchedulingRequirements::getPreviousAllocationId)
+ .filter(Objects::nonNull)
+ .collect(Collectors.toSet());
+ }
+
+ @VisibleForTesting
+ int getNumberOfPendingSlotAssignments() {
+ return pendingSlotAssignments.size();
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionSlotAllocator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionSlotAllocator.java
new file mode 100644
index 0000000..3634f1b
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionSlotAllocator.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.scheduler;
+
+import org.apache.flink.runtime.executiongraph.Execution;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+
+import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Component responsible for assigning slots to a collection of {@link Execution}.
+ */
+public interface ExecutionSlotAllocator {
+
+ /**
+ * Allocate slots for the given executions.
+ *
+ * @param executionVertexSchedulingRequirements The requirements for scheduling the executions.
+ */
+ Collection<SlotExecutionVertexAssignment> allocateSlotsFor(
+ Collection<ExecutionVertexSchedulingRequirements> executionVertexSchedulingRequirements);
+
+ /**
+ * Cancel an ongoing slot request.
+ *
+ * @param executionVertexId identifying which slot request should be canceled.
+ */
+ void cancel(ExecutionVertexID executionVertexId);
+
+ /**
+ * Stop the allocator.
+ */
+ CompletableFuture<Void> stop();
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionVertexSchedulingRequirements.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionVertexSchedulingRequirements.java
new file mode 100644
index 0000000..551719c
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionVertexSchedulingRequirements.java
@@ -0,0 +1,160 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
+import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
+import org.apache.flink.runtime.jobmanager.scheduler.CoLocationConstraint;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+
+import javax.annotation.Nullable;
+
+import java.util.Collection;
+import java.util.Collections;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * The requirements for scheduling a {@link ExecutionVertex}.
+ */
+public class ExecutionVertexSchedulingRequirements {
+
+ private final ExecutionVertexID executionVertexId;
+
+ @Nullable
+ private final AllocationID previousAllocationId;
+
+ private final ResourceProfile resourceProfile;
+
+ @Nullable
+ private final SlotSharingGroupId slotSharingGroupId;
+
+ @Nullable
+ private final CoLocationConstraint coLocationConstraint;
+
+ private final Collection<TaskManagerLocation> preferredLocations;
+
+ private ExecutionVertexSchedulingRequirements(
+ ExecutionVertexID executionVertexId,
+ @Nullable AllocationID previousAllocationId,
+ ResourceProfile resourceProfile,
+ @Nullable SlotSharingGroupId slotSharingGroupId,
+ @Nullable CoLocationConstraint coLocationConstraint,
+ Collection<TaskManagerLocation> preferredLocations) {
+ this.executionVertexId = checkNotNull(executionVertexId);
+ this.previousAllocationId = previousAllocationId;
+ this.resourceProfile = checkNotNull(resourceProfile);
+ this.slotSharingGroupId = slotSharingGroupId;
+ this.coLocationConstraint = coLocationConstraint;
+ this.preferredLocations = checkNotNull(preferredLocations);
+ }
+
+ /**
+ * a {@link ExecutionVertex#MAX_DISTINCT_LOCATIONS_TO_CONSIDER} test.
+ *
+ * @return
+ */
+ public ExecutionVertexID getExecutionVertexId() {
+ return executionVertexId;
+ }
+
+ @Nullable
+ public AllocationID getPreviousAllocationId() {
+ return previousAllocationId;
+ }
+
+ public ResourceProfile getResourceProfile() {
+ return resourceProfile;
+ }
+
+ @Nullable
+ public SlotSharingGroupId getSlotSharingGroupId() {
+ return slotSharingGroupId;
+ }
+
+ @Nullable
+ public CoLocationConstraint getCoLocationConstraint() {
+ return coLocationConstraint;
+ }
+
+ public Collection<TaskManagerLocation> getPreferredLocations() {
+ return preferredLocations;
+ }
+
+ /**
+ * Builder for {@link ExecutionVertexSchedulingRequirements}.
+ */
+ public static class Builder {
+
+ private ExecutionVertexID executionVertexId;
+
+ private AllocationID previousAllocationId;
+
+ private ResourceProfile resourceProfile = ResourceProfile.UNKNOWN;
+
+ private SlotSharingGroupId slotSharingGroupId;
+
+ private CoLocationConstraint coLocationConstraint;
+
+ private Collection<TaskManagerLocation> preferredLocations = Collections.emptyList();
+
+ public Builder withExecutionVertexId(final ExecutionVertexID executionVertexId) {
+ this.executionVertexId = executionVertexId;
+ return this;
+ }
+
+ public Builder withPreviousAllocationId(final AllocationID previousAllocationId) {
+ this.previousAllocationId = previousAllocationId;
+ return this;
+ }
+
+ public Builder withResourceProfile(final ResourceProfile resourceProfile) {
+ this.resourceProfile = resourceProfile;
+ return this;
+ }
+
+ public Builder withSlotSharingGroupId(final SlotSharingGroupId slotSharingGroupId) {
+ this.slotSharingGroupId = slotSharingGroupId;
+ return this;
+ }
+
+ public Builder withCoLocationConstraint(final CoLocationConstraint coLocationConstraint) {
+ this.coLocationConstraint = coLocationConstraint;
+ return this;
+ }
+
+ public Builder withPreferredLocations(final Collection<TaskManagerLocation> preferredLocations) {
+ this.preferredLocations = preferredLocations;
+ return this;
+ }
+
+ public ExecutionVertexSchedulingRequirements build() {
+ return new ExecutionVertexSchedulingRequirements(
+ executionVertexId,
+ previousAllocationId,
+ resourceProfile,
+ slotSharingGroupId,
+ coLocationConstraint,
+ preferredLocations);
+ }
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/InputsLocationsRetriever.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/InputsLocationsRetriever.java
new file mode 100644
index 0000000..c57a5cd
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/InputsLocationsRetriever.java
@@ -0,0 +1,49 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.runtime.executiongraph.Execution;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+
+import java.util.Collection;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Component to retrieve the inputs locations of a {@link Execution}.
+ */
+public interface InputsLocationsRetriever {
+
+ /**
+ * Get the producers of the result partitions consumed by an execution.
+ *
+ * @param executionVertexId identifies the execution
+ * @return the producers of the result partitions group by job vertex id
+ */
+ Collection<Collection<ExecutionVertexID>> getConsumedResultPartitionsProducers(ExecutionVertexID executionVertexId);
+
+ /**
+ * Get the task manager location future for an execution.
+ *
+ * @param executionVertexId identifying the execution
+ * @return the task manager location future
+ */
+ Optional<CompletableFuture<TaskManagerLocation>> getTaskManagerLocation(ExecutionVertexID executionVertexId);
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SlotExecutionVertexAssignment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SlotExecutionVertexAssignment.java
new file mode 100644
index 0000000..b52c6af
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SlotExecutionVertexAssignment.java
@@ -0,0 +1,52 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+import org.apache.flink.runtime.jobmaster.LogicalSlot;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+
+import java.util.concurrent.CompletableFuture;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * The slot assignment for a {@link ExecutionVertex}.
+ */
+public class SlotExecutionVertexAssignment {
+
+ private final ExecutionVertexID executionVertexId;
+
+ private final CompletableFuture<LogicalSlot> logicalSlotFuture;
+
+ public SlotExecutionVertexAssignment(
+ ExecutionVertexID executionVertexId,
+ CompletableFuture<LogicalSlot> logicalSlotFuture) {
+ this.executionVertexId = checkNotNull(executionVertexId);
+ this.logicalSlotFuture = checkNotNull(logicalSlotFuture);
+ }
+
+ public ExecutionVertexID getExecutionVertexId() {
+ return executionVertexId;
+ }
+
+ public CompletableFuture<LogicalSlot> getLogicalSlotFuture() {
+ return logicalSlotFuture;
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorPreferredLocationsTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorPreferredLocationsTest.java
new file mode 100644
index 0000000..d881338
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorPreferredLocationsTest.java
@@ -0,0 +1,116 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.util.TestLogger;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.hasSize;
+import static org.junit.Assert.assertThat;
+
+/**
+ * Tests for the calculation of preferred locations based on inputs in {@link DefaultExecutionSlotAllocator}.
+ */
+public class DefaultExecutionSlotAllocatorPreferredLocationsTest extends TestLogger {
+
+ /**
+ * Tests that the input edge will be ignored if it has too many different locations.
+ */
+ @Test
+ public void testIgnoreEdgeOfTooManyLocations() throws Exception {
+ final ExecutionVertexID consumerId = new ExecutionVertexID(new JobVertexID(), 0);
+ final List<ExecutionVertexID> producerIds = new ArrayList<>(ExecutionVertex.MAX_DISTINCT_LOCATIONS_TO_CONSIDER + 1);
+
+ TestingInputsLocationsRetriever.Builder locationRetrieverBuilder = new TestingInputsLocationsRetriever.Builder();
+ JobVertexID jobVertexID = new JobVertexID();
+ for (int i = 0; i < ExecutionVertex.MAX_DISTINCT_LOCATIONS_TO_CONSIDER + 1; i++) {
+ final ExecutionVertexID producerId = new ExecutionVertexID(jobVertexID, i);
+ locationRetrieverBuilder.connectConsumerToProducer(consumerId, producerId);
+ producerIds.add(producerId);
+ }
+
+ final TestingInputsLocationsRetriever inputsLocationsRetriever = locationRetrieverBuilder.build();
+
+ for (int i = 0; i < ExecutionVertex.MAX_DISTINCT_LOCATIONS_TO_CONSIDER + 1; i++) {
+ inputsLocationsRetriever.markScheduled(producerIds.get(i));
+ }
+
+ CompletableFuture<Collection<TaskManagerLocation>> preferredLocations =
+ DefaultExecutionSlotAllocator.getPreferredLocationsBasedOnInputs(consumerId, inputsLocationsRetriever);
+
+ assertThat(preferredLocations.get(), hasSize(0));
+ }
+
+ /**
+ * Tests that will choose the locations on the edge which has less different number.
+ */
+ @Test
+ public void testChooseLocationsOnEdgeWithLessDifferentNumber() throws Exception {
+ final ExecutionVertexID consumerId = new ExecutionVertexID(new JobVertexID(), 0);
+
+ TestingInputsLocationsRetriever.Builder locationRetrieverBuilder = new TestingInputsLocationsRetriever.Builder();
+ final JobVertexID jobVertexID1 = new JobVertexID();
+ final JobVertexID jobVertexID2 = new JobVertexID();
+ int parallelism1 = 3;
+ int parallelism2 = 5;
+ List<ExecutionVertexID> producers1 = new ArrayList<>(parallelism1);
+ List<ExecutionVertexID> producers2 = new ArrayList<>(parallelism2);
+
+ for (int i = 0; i < parallelism1; i++) {
+ final ExecutionVertexID producerId = new ExecutionVertexID(jobVertexID1, i);
+ producers1.add(producerId);
+ locationRetrieverBuilder.connectConsumerToProducer(consumerId, producerId);
+ }
+
+ for (int i = 0; i < parallelism2; i++) {
+ final ExecutionVertexID producerId = new ExecutionVertexID(jobVertexID2, i);
+ producers2.add(producerId);
+ locationRetrieverBuilder.connectConsumerToProducer(consumerId, producerId);
+ }
+
+ final TestingInputsLocationsRetriever inputsLocationsRetriever = locationRetrieverBuilder.build();
+
+ List<TaskManagerLocation> expectedLocations = new ArrayList<>(parallelism1);
+ for (int i = 0; i < parallelism1; i++) {
+ inputsLocationsRetriever.assignTaskManagerLocation(producers1.get(i));
+ expectedLocations.add(inputsLocationsRetriever.getTaskManagerLocation(producers1.get(i)).get().getNow(null));
+ }
+
+ for (int i = 0; i < parallelism2; i++) {
+ inputsLocationsRetriever.assignTaskManagerLocation(producers2.get(i));
+ }
+
+ CompletableFuture<Collection<TaskManagerLocation>> preferredLocations =
+ DefaultExecutionSlotAllocator.getPreferredLocationsBasedOnInputs(consumerId, inputsLocationsRetriever);
+
+ assertThat(preferredLocations.get(), containsInAnyOrder(expectedLocations.toArray()));
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorTest.java
new file mode 100644
index 0000000..4c0dd01
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocatorTest.java
@@ -0,0 +1,336 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.api.common.time.Time;
+import org.apache.flink.api.java.tuple.Tuple3;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
+import org.apache.flink.runtime.clusterframework.types.SlotProfile;
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.jobmanager.scheduler.CoLocationConstraint;
+import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
+import org.apache.flink.runtime.jobmanager.scheduler.ScheduledUnit;
+import org.apache.flink.runtime.jobmaster.LogicalSlot;
+import org.apache.flink.runtime.jobmaster.SlotRequestId;
+import org.apache.flink.runtime.jobmaster.TestingLogicalSlot;
+import org.apache.flink.runtime.jobmaster.slotpool.SlotProvider;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.util.TestLogger;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import javax.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.hasSize;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * Tests for {@link DefaultExecutionSlotAllocator}.
+ */
+public class DefaultExecutionSlotAllocatorTest extends TestLogger {
+
+ private AllocationToggableSlotProvider slotProvider;
+
+ @Before
+ public void setUp() throws Exception {
+ slotProvider = new AllocationToggableSlotProvider();
+ }
+
+ /**
+ * Tests that consumers will get slots after producers are fulfilled.
+ */
+ @Test
+ public void testConsumersAssignedToSlotsAfterProducers() {
+ final ExecutionVertexID producerId = new ExecutionVertexID(new JobVertexID(), 0);
+ final ExecutionVertexID consumerId = new ExecutionVertexID(new JobVertexID(), 0);
+
+ final TestingInputsLocationsRetriever inputsLocationsRetriever = new TestingInputsLocationsRetriever.Builder()
+ .connectConsumerToProducer(consumerId, producerId)
+ .build();
+
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator(inputsLocationsRetriever);
+
+ inputsLocationsRetriever.markScheduled(producerId);
+ inputsLocationsRetriever.markScheduled(consumerId);
+
+ final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = createSchedulingRequirements(producerId, consumerId);
+ final Collection<SlotExecutionVertexAssignment> slotExecutionVertexAssignments = executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
+ assertThat(slotExecutionVertexAssignments, hasSize(2));
+
+ final SlotExecutionVertexAssignment producerSlotAssignment = findSlotAssignmentByExecutionVertexId(producerId, slotExecutionVertexAssignments);
+ final SlotExecutionVertexAssignment consumerSlotAssignment = findSlotAssignmentByExecutionVertexId(consumerId, slotExecutionVertexAssignments);
+
+ assertTrue(producerSlotAssignment.getLogicalSlotFuture().isDone());
+ assertFalse(consumerSlotAssignment.getLogicalSlotFuture().isDone());
+
+ inputsLocationsRetriever.assignTaskManagerLocation(producerId);
+
+ assertTrue(consumerSlotAssignment.getLogicalSlotFuture().isDone());
+ assertEquals(0, executionSlotAllocator.getNumberOfPendingSlotAssignments());
+ }
+
+ /**
+ * Tests that validate the parameters when calling allocateSlot in SlotProvider.
+ */
+ @Test
+ public void testAllocateSlotsParameters() {
+ final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
+ final AllocationID allocationId = new AllocationID();
+ final SlotSharingGroupId sharingGroupId = new SlotSharingGroupId();
+ final ResourceProfile resourceProfile = new ResourceProfile(0.5, 250);
+ final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
+ final Collection<TaskManagerLocation> taskManagerLocations = Collections.singleton(new LocalTaskManagerLocation());
+
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();
+
+ final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
+ new ExecutionVertexSchedulingRequirements.Builder()
+ .withExecutionVertexId(executionVertexId)
+ .withPreviousAllocationId(allocationId)
+ .withSlotSharingGroupId(sharingGroupId)
+ .withPreferredLocations(taskManagerLocations)
+ .withResourceProfile(resourceProfile)
+ .withCoLocationConstraint(coLocationConstraint)
+ .build()
+ );
+
+ executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
+ assertThat(slotProvider.getSlotAllocationRequests(), hasSize(1));
+
+ ScheduledUnit expectedTask = slotProvider.getSlotAllocationRequests().get(0).f1;
+ SlotProfile expectedSlotProfile = slotProvider.getSlotAllocationRequests().get(0).f2;
+
+ assertEquals(sharingGroupId, expectedTask.getSlotSharingGroupId());
+ assertEquals(coLocationConstraint, expectedTask.getCoLocationConstraint());
+ assertThat(expectedSlotProfile.getPreferredAllocations(), contains(allocationId));
+ assertThat(expectedSlotProfile.getPreviousExecutionGraphAllocations(), contains(allocationId));
+ assertEquals(resourceProfile, expectedSlotProfile.getResourceProfile());
+ assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocations.toArray()));
+ }
+
+ /**
+ * Tests that cancels an execution vertex which is not existed.
+ */
+ @Test
+ public void testCancelNonExistingExecutionVertex() {
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();
+
+ ExecutionVertexID inValidExecutionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
+ executionSlotAllocator.cancel(inValidExecutionVertexId);
+
+ assertThat(slotProvider.getCancelledSlotRequestIds(), is(empty()));
+ }
+
+ /**
+ * Tests that cancels a slot request which has already been fulfilled.
+ */
+ @Test
+ public void testCancelFulfilledSlotRequest() {
+ final ExecutionVertexID producerId = new ExecutionVertexID(new JobVertexID(), 0);
+
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();
+
+ final List<ExecutionVertexSchedulingRequirements> schedulingRequirements =
+ createSchedulingRequirements(producerId);
+ executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
+
+ executionSlotAllocator.cancel(producerId);
+
+ assertThat(slotProvider.getCancelledSlotRequestIds(), is(empty()));
+ }
+
+ /**
+ * Tests that cancels a slot request which has not been fulfilled.
+ */
+ @Test
+ public void testCancelUnFulfilledSlotRequest() throws Exception {
+ final ExecutionVertexID producerId = new ExecutionVertexID(new JobVertexID(), 0);
+
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();
+
+ slotProvider.disableSlotAllocation();
+ final List<ExecutionVertexSchedulingRequirements> schedulingRequirements =
+ createSchedulingRequirements(producerId);
+ Collection<SlotExecutionVertexAssignment> assignments = executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
+
+ executionSlotAllocator.cancel(producerId);
+
+ assertThat(slotProvider.getCancelledSlotRequestIds(), hasSize(1));
+ assertThat(slotProvider.getCancelledSlotRequestIds(), contains(slotProvider.getReceivedSlotRequestIds().toArray()));
+
+ try {
+ assignments.iterator().next().getLogicalSlotFuture().get();
+ fail("Expect a CancellationException but got nothing.");
+ } catch (CancellationException ignored) {
+ // Expected exception
+ }
+ }
+
+ /**
+ * Tests that all unfulfilled slot requests will be cancelled when stopped.
+ */
+ @Test
+ public void testStop() throws Exception {
+ final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
+
+ final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();
+
+ slotProvider.disableSlotAllocation();
+ final List<ExecutionVertexSchedulingRequirements> schedulingRequirements =
+ createSchedulingRequirements(executionVertexId);
+ executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
+
+ executionSlotAllocator.stop().get();
+
+ assertThat(slotProvider.getCancelledSlotRequestIds(), hasSize(1));
+ assertThat(slotProvider.getCancelledSlotRequestIds(), contains(slotProvider.getReceivedSlotRequestIds().toArray()));
+ assertEquals(0, executionSlotAllocator.getNumberOfPendingSlotAssignments());
+ }
+
+ /**
+ * Tests that all prior allocation ids are computed by union all previous allocation ids in scheduling requirements.
+ */
+ @Test
+ public void testComputeAllPriorAllocationIds() {
+ List<AllocationID> expectAllocationIds = Arrays.asList(new AllocationID(), new AllocationID());
+ List<ExecutionVertexSchedulingRequirements> testSchedulingRequirements = Arrays.asList(
+ new ExecutionVertexSchedulingRequirements.Builder().
+ withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 0)).
+ withPreviousAllocationId(expectAllocationIds.get(0)).
+ build(),
+ new ExecutionVertexSchedulingRequirements.Builder().
+ withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 1)).
+ withPreviousAllocationId(expectAllocationIds.get(0)).
+ build(),
+ new ExecutionVertexSchedulingRequirements.Builder().
+ withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 2)).
+ withPreviousAllocationId(expectAllocationIds.get(1)).
+ build(),
+ new ExecutionVertexSchedulingRequirements.Builder().
+ withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 3)).
+ build()
+ );
+
+ Set<AllocationID> allPriorAllocationIds = DefaultExecutionSlotAllocator.computeAllPriorAllocationIds(testSchedulingRequirements);
+ assertThat(allPriorAllocationIds, containsInAnyOrder(expectAllocationIds.toArray()));
+ }
+
+ private DefaultExecutionSlotAllocator createExecutionSlotAllocator() {
+ return createExecutionSlotAllocator(new TestingInputsLocationsRetriever.Builder().build());
+ }
+
+ private DefaultExecutionSlotAllocator createExecutionSlotAllocator(InputsLocationsRetriever inputsLocationsRetriever) {
+ return new DefaultExecutionSlotAllocator(slotProvider, inputsLocationsRetriever, Time.seconds(10));
+ }
+
+ private List<ExecutionVertexSchedulingRequirements> createSchedulingRequirements(ExecutionVertexID... executionVertexIds) {
+ List<ExecutionVertexSchedulingRequirements> schedulingRequirements = new ArrayList<>(executionVertexIds.length);
+
+ for (ExecutionVertexID executionVertexId : executionVertexIds) {
+ schedulingRequirements.add(new ExecutionVertexSchedulingRequirements.Builder()
+ .withExecutionVertexId(executionVertexId).build());
+ }
+ return schedulingRequirements;
+ }
+
+ private SlotExecutionVertexAssignment findSlotAssignmentByExecutionVertexId(
+ ExecutionVertexID executionVertexId,
+ Collection<SlotExecutionVertexAssignment> slotExecutionVertexAssignments) {
+ return slotExecutionVertexAssignments.stream()
+ .filter(slotExecutionVertexAssignment -> slotExecutionVertexAssignment.getExecutionVertexId().equals(executionVertexId))
+ .findFirst()
+ .orElseThrow(() -> new IllegalArgumentException(String.format(
+ "SlotExecutionVertexAssignment with execution vertex id %s not found",
+ executionVertexId)));
+ }
+
+ private static class AllocationToggableSlotProvider implements SlotProvider {
+
+ private final List<Tuple3<SlotRequestId, ScheduledUnit, SlotProfile>> slotAllocationRequests = new ArrayList<>();
+
+ private final List<SlotRequestId> cancelledSlotRequestIds = new ArrayList<>();
+
+ private boolean slotAllocationDisabled;
+
+ @Override
+ public CompletableFuture<LogicalSlot> allocateSlot(
+ SlotRequestId slotRequestId,
+ ScheduledUnit task,
+ SlotProfile slotProfile,
+ boolean allowQueued,
+ Time timeout) {
+
+ slotAllocationRequests.add(Tuple3.of(slotRequestId, task, slotProfile));
+ if (slotAllocationDisabled) {
+ return new CompletableFuture<>();
+ } else {
+ return CompletableFuture.completedFuture(new TestingLogicalSlot());
+ }
+ }
+
+ @Override
+ public void cancelSlotRequest(
+ final SlotRequestId slotRequestId,
+ @Nullable final SlotSharingGroupId slotSharingGroupId,
+ final Throwable cause) {
+ cancelledSlotRequestIds.add(slotRequestId);
+ }
+
+ public List<Tuple3<SlotRequestId, ScheduledUnit, SlotProfile>> getSlotAllocationRequests() {
+ return Collections.unmodifiableList(slotAllocationRequests);
+ }
+
+ public List<SlotRequestId> getReceivedSlotRequestIds() {
+ return slotAllocationRequests.stream()
+ .map(requestTuple -> requestTuple.f0)
+ .collect(Collectors.toList());
+ }
+
+ public List<SlotRequestId> getCancelledSlotRequestIds() {
+ return Collections.unmodifiableList(cancelledSlotRequestIds);
+ }
+
+ public void disableSlotAllocation() {
+ slotAllocationDisabled = true;
+ }
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingInputsLocationsRetriever.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingInputsLocationsRetriever.java
new file mode 100644
index 0000000..113740a
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingInputsLocationsRetriever.java
@@ -0,0 +1,94 @@
+/*
+ * 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.scheduler;
+
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+/**
+ * A simple inputs locations retriever for testing purposes.
+ */
+class TestingInputsLocationsRetriever implements InputsLocationsRetriever {
+
+ private final Map<ExecutionVertexID, List<ExecutionVertexID>> producersByConsumer;
+
+ private final Map<ExecutionVertexID, CompletableFuture<TaskManagerLocation>> taskManagerLocationsByVertex = new HashMap<>();
+
+ TestingInputsLocationsRetriever(final Map<ExecutionVertexID, List<ExecutionVertexID>> producersByConsumer) {
+ this.producersByConsumer = new HashMap<>(producersByConsumer);
+ }
+
+ @Override
+ public Collection<Collection<ExecutionVertexID>> getConsumedResultPartitionsProducers(final ExecutionVertexID executionVertexId) {
+ final Map<JobVertexID, List<ExecutionVertexID>> executionVerticesByJobVertex =
+ producersByConsumer.getOrDefault(executionVertexId, Collections.emptyList())
+ .stream()
+ .collect(Collectors.groupingBy(ExecutionVertexID::getJobVertexId));
+
+ return new ArrayList<>(executionVerticesByJobVertex.values());
+ }
+
+ @Override
+ public Optional<CompletableFuture<TaskManagerLocation>> getTaskManagerLocation(final ExecutionVertexID executionVertexId) {
+ return Optional.ofNullable(taskManagerLocationsByVertex.get(executionVertexId));
+ }
+
+ public void markScheduled(final ExecutionVertexID executionVertexId) {
+ taskManagerLocationsByVertex.put(executionVertexId, new CompletableFuture<>());
+ }
+
+ public void assignTaskManagerLocation(final ExecutionVertexID executionVertexId) {
+ taskManagerLocationsByVertex.compute(executionVertexId, (key, future) -> {
+ if (future == null) {
+ return CompletableFuture.completedFuture(new LocalTaskManagerLocation());
+ }
+ future.complete(new LocalTaskManagerLocation());
+ return future;
+ });
+ }
+
+ static class Builder {
+
+ private final Map<ExecutionVertexID, List<ExecutionVertexID>> producersByConsumer = new HashMap<>();
+
+ public Builder connectConsumerToProducer(final ExecutionVertexID consumer, final ExecutionVertexID producer) {
+ producersByConsumer
+ .computeIfAbsent(consumer, (key) -> new ArrayList<>())
+ .add(producer);
+ return this;
+ }
+
+ public TestingInputsLocationsRetriever build() {
+ return new TestingInputsLocationsRetriever(producersByConsumer);
+ }
+
+ }
+}