@Test public void testAnyPreferredLocationCalculation() throws ExecutionException, InterruptedException { final TaskManagerLocation taskManagerLocation1 = new LocalTaskManagerLocation(); final TaskManagerLocation taskManagerLocation3 = new LocalTaskManagerLocation(); final CompletableFuture<TaskManagerLocation> locationFuture1 = CompletableFuture.completedFuture(taskManagerLocation1); final CompletableFuture<TaskManagerLocation> locationFuture2 = new CompletableFuture<>(); final CompletableFuture<TaskManagerLocation> locationFuture3 = CompletableFuture.completedFuture(taskManagerLocation3); final Execution execution = SchedulerTestUtils.getTestVertex(Arrays.asList(locationFuture1, locationFuture2, locationFuture3)); CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = execution.calculatePreferredLocations(LocationPreferenceConstraint.ANY); assertTrue(preferredLocationsFuture.isDone()); final Collection<TaskManagerLocation> preferredLocations = preferredLocationsFuture.get(); assertThat(preferredLocations, containsInAnyOrder(taskManagerLocation1, taskManagerLocation3)); }
@VisibleForTesting public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) { final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations(); final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture; switch(locationPreferenceConstraint) { case ALL: preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures); break; case ANY: final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size()); for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) { if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) { final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null); if (taskManagerLocation == null) { throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug."); } completedTaskManagerLocations.add(taskManagerLocation); } } preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations); break; default: throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.'); } return preferredLocationsFuture; }
Execution implements AccessExecution, Archiveable<ArchivedExecution>, LogicalSlot.Payload { @VisibleForTesting public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) { final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations(); final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture; switch(locationPreferenceConstraint) { case ALL: preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures); break; case ANY: final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size()); for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) { if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) { final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null); if (taskManagerLocation == null) { throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug."); } completedTaskManagerLocations.add(taskManagerLocation); } } preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations); break; default: throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.'); } return preferredLocationsFuture; } }
Execution implements AccessExecution, Archiveable<ArchivedExecution>, LogicalSlot.Payload { @VisibleForTesting public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) { final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations(); final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture; switch(locationPreferenceConstraint) { case ALL: preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures); break; case ANY: final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size()); for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) { if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) { final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null); if (taskManagerLocation == null) { throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug."); } completedTaskManagerLocations.add(taskManagerLocation); } } preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations); break; default: throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.'); } return preferredLocationsFuture; } Execution(
			Executor executor,
			ExecutionVertex vertex,
			int attemptNumber,
			long globalModVersion,
			long startTimestamp,
			Time rpcTimeout); }
Execution implements AccessExecution, Archiveable<ArchivedExecution>, LogicalSlot.Payload { @VisibleForTesting public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) { final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations(); final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture; switch(locationPreferenceConstraint) { case ALL: preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures); break; case ANY: final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size()); for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) { if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) { final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null); if (taskManagerLocation == null) { throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug."); } completedTaskManagerLocations.add(taskManagerLocation); } } preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations); break; default: throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.'); } return preferredLocationsFuture; } Execution(
			Executor executor,
			ExecutionVertex vertex,
			int attemptNumber,
			long globalModVersion,
			long startTimestamp,
			Time rpcTimeout); ExecutionVertex getVertex(); @Override ExecutionAttemptID getAttemptId(); @Override int getAttemptNumber(); @Override ExecutionState getState(); @Nullable AllocationID getAssignedAllocationID(); long getGlobalModVersion(); CompletableFuture<TaskManagerLocation> getTaskManagerLocationFuture(); LogicalSlot getAssignedResource(); @Override TaskManagerLocation getAssignedResourceLocation(); Throwable getFailureCause(); @Override String getFailureCauseAsString(); @Override long[] getStateTimestamps(); @Override long getStateTimestamp(ExecutionState state); boolean isFinished(); @Nullable JobManagerTaskRestore getTaskRestore(); void setInitialState(@Nullable JobManagerTaskRestore taskRestore); @Override CompletableFuture<ExecutionState> getTerminalStateFuture(); CompletableFuture<?> getReleaseFuture(); CompletableFuture<Void> scheduleForExecution(); CompletableFuture<Void> scheduleForExecution(
			SlotProvider slotProvider,
			boolean queued,
			LocationPreferenceConstraint locationPreferenceConstraint,
			@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds); CompletableFuture<Execution> allocateAndAssignSlotForExecution(
			SlotProvider slotProvider,
			boolean queued,
			LocationPreferenceConstraint locationPreferenceConstraint,
			@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds,
			Time allocationTimeout); void deploy(); void stop(); void cancel(); CompletableFuture<?> suspend(); @Override void fail(Throwable t); CompletableFuture<StackTraceSampleResponse> requestStackTraceSample(
			int sampleId,
			int numSamples,
			Time delayBetweenSamples,
			int maxStackTraceDepth,
			Time timeout); void notifyCheckpointComplete(long checkpointId, long timestamp); void triggerCheckpoint(long checkpointId, long timestamp, CheckpointOptions checkpointOptions); @VisibleForTesting CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint); String getVertexWithAttempt(); void setAccumulators(Map<String, Accumulator<?, ?>> userAccumulators); Map<String, Accumulator<?, ?>> getUserAccumulators(); @Override StringifiedAccumulatorResult[] getUserAccumulatorsStringified(); @Override int getParallelSubtaskIndex(); @Override IOMetrics getIOMetrics(); @Override String toString(); @Override ArchivedExecution archive(); }
Execution implements AccessExecution, Archiveable<ArchivedExecution>, LogicalSlot.Payload { @VisibleForTesting public CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint) { final Collection<CompletableFuture<TaskManagerLocation>> preferredLocationFutures = getVertex().getPreferredLocations(); final CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture; switch(locationPreferenceConstraint) { case ALL: preferredLocationsFuture = FutureUtils.combineAll(preferredLocationFutures); break; case ANY: final ArrayList<TaskManagerLocation> completedTaskManagerLocations = new ArrayList<>(preferredLocationFutures.size()); for (CompletableFuture<TaskManagerLocation> preferredLocationFuture : preferredLocationFutures) { if (preferredLocationFuture.isDone() && !preferredLocationFuture.isCompletedExceptionally()) { final TaskManagerLocation taskManagerLocation = preferredLocationFuture.getNow(null); if (taskManagerLocation == null) { throw new FlinkRuntimeException("TaskManagerLocationFuture was completed with null. This indicates a programming bug."); } completedTaskManagerLocations.add(taskManagerLocation); } } preferredLocationsFuture = CompletableFuture.completedFuture(completedTaskManagerLocations); break; default: throw new RuntimeException("Unknown LocationPreferenceConstraint " + locationPreferenceConstraint + '.'); } return preferredLocationsFuture; } Execution(
			Executor executor,
			ExecutionVertex vertex,
			int attemptNumber,
			long globalModVersion,
			long startTimestamp,
			Time rpcTimeout); ExecutionVertex getVertex(); @Override ExecutionAttemptID getAttemptId(); @Override int getAttemptNumber(); @Override ExecutionState getState(); @Nullable AllocationID getAssignedAllocationID(); long getGlobalModVersion(); CompletableFuture<TaskManagerLocation> getTaskManagerLocationFuture(); LogicalSlot getAssignedResource(); @Override TaskManagerLocation getAssignedResourceLocation(); Throwable getFailureCause(); @Override String getFailureCauseAsString(); @Override long[] getStateTimestamps(); @Override long getStateTimestamp(ExecutionState state); boolean isFinished(); @Nullable JobManagerTaskRestore getTaskRestore(); void setInitialState(@Nullable JobManagerTaskRestore taskRestore); @Override CompletableFuture<ExecutionState> getTerminalStateFuture(); CompletableFuture<?> getReleaseFuture(); CompletableFuture<Void> scheduleForExecution(); CompletableFuture<Void> scheduleForExecution(
			SlotProvider slotProvider,
			boolean queued,
			LocationPreferenceConstraint locationPreferenceConstraint,
			@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds); CompletableFuture<Execution> allocateAndAssignSlotForExecution(
			SlotProvider slotProvider,
			boolean queued,
			LocationPreferenceConstraint locationPreferenceConstraint,
			@Nonnull Set<AllocationID> allPreviousExecutionGraphAllocationIds,
			Time allocationTimeout); void deploy(); void stop(); void cancel(); CompletableFuture<?> suspend(); @Override void fail(Throwable t); CompletableFuture<StackTraceSampleResponse> requestStackTraceSample(
			int sampleId,
			int numSamples,
			Time delayBetweenSamples,
			int maxStackTraceDepth,
			Time timeout); void notifyCheckpointComplete(long checkpointId, long timestamp); void triggerCheckpoint(long checkpointId, long timestamp, CheckpointOptions checkpointOptions); @VisibleForTesting CompletableFuture<Collection<TaskManagerLocation>> calculatePreferredLocations(LocationPreferenceConstraint locationPreferenceConstraint); String getVertexWithAttempt(); void setAccumulators(Map<String, Accumulator<?, ?>> userAccumulators); Map<String, Accumulator<?, ?>> getUserAccumulators(); @Override StringifiedAccumulatorResult[] getUserAccumulatorsStringified(); @Override int getParallelSubtaskIndex(); @Override IOMetrics getIOMetrics(); @Override String toString(); @Override ArchivedExecution archive();  }