[FLINK-16180][runtime] Replace the nullable vertexExecution in Schedu…

…ledUnit with a non-null executionVertexId
mbalassi · Feb 28, 2020 · 9a3eb5b · 9a3eb5b
1 parent 8f6fcd2
commit 9a3eb5b
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 31 deletions.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/scheduler/ScheduledUnit.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/scheduler/ScheduledUnit.java
@@ -18,23 +18,21 @@
 
 package org.apache.flink.runtime.jobmanager.scheduler;
 
+import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.runtime.executiongraph.Execution;
 import org.apache.flink.runtime.instance.SlotSharingGroupId;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
 import org.apache.flink.util.Preconditions;
 
 import javax.annotation.Nullable;
 
 /**
- * ScheduledUnit contains the information necessary to allocate a slot for the given
- * {@link JobVertexID}.
+ * ScheduledUnit contains the information necessary to allocate a slot for the given task.
  */
 public class ScheduledUnit {
 
- @Nullable
- private final Execution vertexExecution;
-
- private final JobVertexID jobVertexId;
+ private final ExecutionVertexID executionVertexId;
 
  @Nullable
  private final SlotSharingGroupId slotSharingGroupId;
@@ -44,18 +42,16 @@ public class ScheduledUnit {
 
  // --------------------------------------------------------------------------------------------
 
+ @VisibleForTesting
  public ScheduledUnit(Execution task) {
  this(
- Preconditions.checkNotNull(task),
- task.getVertex().getJobvertexId(),
- null,
+ task,
  null);
  }
 
  public ScheduledUnit(Execution task, @Nullable SlotSharingGroupId slotSharingGroupId) {
  this(
- Preconditions.checkNotNull(task),
- task.getVertex().getJobvertexId(),
+ task,
  slotSharingGroupId,
  null);
  }
@@ -65,31 +61,28 @@ public ScheduledUnit(
  @Nullable SlotSharingGroupId slotSharingGroupId,
  @Nullable CoLocationConstraint coLocationConstraint) {
  this(
- Preconditions.checkNotNull(task),
- task.getVertex().getJobvertexId(),
+ Preconditions.checkNotNull(task).getVertex().getID(),
  slotSharingGroupId,
  coLocationConstraint);
  }
 
+ @VisibleForTesting
  public ScheduledUnit(
  JobVertexID jobVertexId,
  @Nullable SlotSharingGroupId slotSharingGroupId,
  @Nullable CoLocationConstraint coLocationConstraint) {
  this(
- null,
- jobVertexId,
+ new ExecutionVertexID(jobVertexId, 0),
  slotSharingGroupId,
  coLocationConstraint);
  }
 
  public ScheduledUnit(
- @Nullable Execution task,
- JobVertexID jobVertexId,
- @Nullable SlotSharingGroupId slotSharingGroupId,
- @Nullable CoLocationConstraint coLocationConstraint) {
+ ExecutionVertexID executionVertexId,
+ @Nullable SlotSharingGroupId slotSharingGroupId,
+ @Nullable CoLocationConstraint coLocationConstraint) {
 
- this.vertexExecution = task;
- this.jobVertexId = Preconditions.checkNotNull(jobVertexId);
+ this.executionVertexId = Preconditions.checkNotNull(executionVertexId);
  this.slotSharingGroupId = slotSharingGroupId;
  this.coLocationConstraint = coLocationConstraint;
 
@@ -98,12 +91,11 @@ public ScheduledUnit(
  // --------------------------------------------------------------------------------------------
 
  public JobVertexID getJobVertexId() {
- return jobVertexId;
+ return executionVertexId.getJobVertexId();
  }
 
- @Nullable
- public Execution getTaskToExecute() {
- return vertexExecution;
+ public int getSubtaskIndex() {
+ return executionVertexId.getSubtaskIndex();
  }
 
  @Nullable
@@ -120,7 +112,7 @@ public CoLocationConstraint getCoLocationConstraint() {
 
  @Override
  public String toString() {
- return "{task=" + vertexExecution.getVertexWithAttempt() + ", sharingUnit=" + slotSharingGroupId +
+ return "{task=" + executionVertexId + ", sharingUnit=" + slotSharingGroupId +
  ", locationConstraint=" + coLocationConstraint + '}';
  }
 }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SchedulerImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SchedulerImpl.java
@@ -135,7 +135,7 @@ private CompletableFuture<LogicalSlot> allocateSlotInternal(
  ScheduledUnit scheduledUnit,
  SlotProfile slotProfile,
  @Nullable Time allocationTimeout) {
- log.debug("Received slot request [{}] for task: {}", slotRequestId, scheduledUnit.getTaskToExecute());
+ log.debug("Received slot request [{}] for task: {}", slotRequestId, scheduledUnit.getJobVertexId());
 
  componentMainThreadExecutor.assertRunningInMainThread();
 

diff --git a/...ntime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java b/...ntime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
@@ -104,7 +104,7 @@ public List<SlotExecutionVertexAssignment> allocateSlotsFor(
  slotProviderStrategy.allocateSlot(
  slotRequestId,
  new ScheduledUnit(
- executionVertexId.getJobVertexId(),
+ executionVertexId,
  slotSharingGroupId,
  schedulingRequirements.getCoLocationConstraint()),
  SlotProfile.priorAllocation(

diff --git a/...runtime/src/test/java/org/apache/flink/runtime/executiongraph/ProgrammedSlotProvider.java b/...runtime/src/test/java/org/apache/flink/runtime/executiongraph/ProgrammedSlotProvider.java
@@ -118,8 +118,8 @@ public CompletableFuture<LogicalSlot> allocateSlot(
  ScheduledUnit task,
  SlotProfile slotProfile,
  Time allocationTimeout) {
- JobVertexID vertexId = task.getTaskToExecute().getVertex().getJobvertexId();
- int subtask = task.getTaskToExecute().getParallelSubtaskIndex();
+ final JobVertexID vertexId = task.getJobVertexId();
+ final int subtask = task.getSubtaskIndex();
 
  CompletableFuture<LogicalSlot>[] forTask = slotFutures.get(vertexId);
  if (forTask != null) {

diff --git a/...ntime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/DummyScheduledUnit.java b/...ntime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/DummyScheduledUnit.java
@@ -26,7 +26,6 @@
 public class DummyScheduledUnit extends ScheduledUnit {
  public DummyScheduledUnit() {
  super(
- null,
  new JobVertexID(),
  null,
  null);