org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation Java Examples

@Test
public void testCancelWhileRestarting() throws Exception {
	// We want to manually control the restart and delay
	try (SlotPool slotPool = createSlotPoolImpl()) {
		TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
		final ExecutionGraph executionGraph = TestingExecutionGraphBuilder.newBuilder()
			.setRestartStrategy(new InfiniteDelayRestartStrategy())
			.setTaskManagerLocation(taskManagerLocation)
			.buildAndScheduleForExecution(slotPool);

		// Release the TaskManager and wait for the job to restart
		slotPool.releaseTaskManager(taskManagerLocation.getResourceID(), new Exception("Test Exception"));
		assertEquals(JobStatus.RESTARTING, executionGraph.getState());

		// Canceling needs to abort the restart
		executionGraph.cancel();

		assertEquals(JobStatus.CANCELED, executionGraph.getState());

		// The restart has been aborted
		executionGraph.restart(executionGraph.getGlobalModVersion());

		assertEquals(JobStatus.CANCELED, executionGraph.getState());
	}

}
 

@Nonnull
private ProgrammedSlotProvider createProgrammedSlotProvider(
	int parallelism,
	Collection<JobVertexID> jobVertexIds,
	SlotOwner slotOwner) {
	final ProgrammedSlotProvider slotProvider = new ProgrammedSlotProvider(parallelism);

	for (JobVertexID jobVertexId : jobVertexIds) {
		for (int i = 0; i < parallelism; i++) {
			final SimpleSlot slot = new SimpleSlot(
				slotOwner,
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway(),
				null,
				null);

			slotProvider.addSlot(jobVertexId, 0, CompletableFuture.completedFuture(slot));
		}
	}

	return slotProvider;
}
 

public TaskManagerServicesBuilder() {
	taskManagerLocation = new LocalTaskManagerLocation();
	memoryManager = new MemoryManager(
		MemoryManager.MIN_PAGE_SIZE,
		1,
		MemoryManager.MIN_PAGE_SIZE,
		MemoryType.HEAP,
		false);
	ioManager = mock(IOManager.class);
	networkEnvironment = mock(NetworkEnvironment.class);
	broadcastVariableManager = new BroadcastVariableManager();
	taskSlotTable = mock(TaskSlotTable.class);
	jobManagerTable = new JobManagerTable();
	jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());
	taskStateManager = mock(TaskExecutorLocalStateStoresManager.class);
}
 

/**
 * Tests that it is possible to fail a graph via a call to
 * {@link ExecutionGraph#failGlobal(Throwable)} after cancellation.
 */
@Test
public void testFailExecutionGraphAfterCancel() throws Exception {
	try (SlotPool slotPool = createSlotPoolImpl()) {
		ExecutionGraph eg = TestingExecutionGraphBuilder
			.newBuilder()
			.setRestartStrategy(new InfiniteDelayRestartStrategy())
			.setJobGraph(createJobGraphToCancel())
			.setSlotProvider(createSchedulerWithSlots(slotPool, new LocalTaskManagerLocation(), 2))
			.build();

		startAndScheduleExecutionGraph(eg);

		// Fail right after cancel (for example with concurrent slot release)
		eg.cancel();
		assertEquals(JobStatus.CANCELLING, eg.getState());

		eg.failGlobal(new Exception("Test Exception"));
		assertEquals(JobStatus.FAILING, eg.getState());

		Execution execution = eg.getAllExecutionVertices().iterator().next().getCurrentExecutionAttempt();

		execution.completeCancelling();
		assertEquals(JobStatus.RESTARTING, eg.getState());
	}
}
 

@Before
public void setup() throws IOException {
	rpc = new TestingRpcService();
	timerService = new TimerService<>(TestingUtils.defaultExecutor(), timeout.toMilliseconds());

	dummyBlobCacheService = new BlobCacheService(
		new Configuration(),
		new VoidBlobStore(),
		null);

	configuration = new Configuration();
	taskManagerConfiguration = TaskManagerConfiguration.fromConfiguration(configuration);

	taskManagerLocation = new LocalTaskManagerLocation();
	jobId = new JobID();

	testingFatalErrorHandler = new TestingFatalErrorHandler();

	haServices = new TestingHighAvailabilityServices();
	resourceManagerLeaderRetriever = new SettableLeaderRetrievalService();
	jobManagerLeaderRetriever = new SettableLeaderRetrievalService();
	haServices.setResourceManagerLeaderRetriever(resourceManagerLeaderRetriever);
	haServices.setJobMasterLeaderRetriever(jobId, jobManagerLeaderRetriever);
}
 

/**
 * Tests that any preferred locations are calculated.
 */
@Test
public void testAnyPreferredLocationCalculation() throws Exception {
	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 = getExecution(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));
}
 

/**
 * Tests that the location preferences are honoured when looking for a resolved slot.
 */
@Test
public void testGetResolvedSlotWithLocationPreferences() {
	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot rootSlot1 = createRootSlot(new LocalTaskManagerLocation(), slotSharingManager);

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	SlotSharingManager.MultiTaskSlot rootSlot2 = createRootSlot(taskManagerLocation, slotSharingManager);

	AbstractID groupId = new AbstractID();

	SlotProfile slotProfile = SlotProfile.preferredLocality(ResourceProfile.UNKNOWN, Collections.singleton(taskManagerLocation));

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	final LocationPreferenceSlotSelectionStrategy locationPreferenceSlotSelectionStrategy = LocationPreferenceSlotSelectionStrategy.createDefault();
	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		locationPreferenceSlotSelectionStrategy.selectBestSlotForProfile(slotInfos, slotProfile).get();
	SlotSharingManager.MultiTaskSlot resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());

	assertNotNull(resolvedRootSlot);
	assertEquals(Locality.LOCAL, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot2.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());

	// occupy the slot
	resolvedRootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId,
		slotInfoAndLocality.getLocality());

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	slotInfoAndLocality = locationPreferenceSlotSelectionStrategy.selectBestSlotForProfile(slotInfos, slotProfile).get();
	resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());
	assertNotNull(resolvedRootSlot);
	assertNotSame(Locality.LOCAL, (slotInfoAndLocality.getLocality()));
	assertEquals(rootSlot1.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());
}
 

/**
 * Tests that a suspend call while restarting a job, will abort the restarting.
 */
@Test
public void testSuspendWhileRestarting() throws Exception {
	TestRestartStrategy controllableRestartStrategy = TestRestartStrategy.manuallyTriggered();
	try (SlotPool slotPool = createSlotPoolImpl()) {
		TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
		ExecutionGraph eg = TestingExecutionGraphBuilder
			.newBuilder()
			.setJobGraph(createJobGraph())
			.setRestartStrategy(controllableRestartStrategy)
			.setSlotProvider(createSchedulerWithSlots(slotPool, taskManagerLocation))
			.build();

		startAndScheduleExecutionGraph(eg);

		// Release the TaskManager and wait for the job to restart
		slotPool.releaseTaskManager(taskManagerLocation.getResourceID(), new Exception("Test Exception"));

		assertEquals(1, controllableRestartStrategy.getNumberOfQueuedActions());

		assertEquals(JobStatus.RESTARTING, eg.getState());

		eg.suspend(new Exception("Test exception"));

		assertEquals(JobStatus.SUSPENDED, eg.getState());

		controllableRestartStrategy.triggerAll().join();

		assertEquals(JobStatus.SUSPENDED, eg.getState());
	}
}
 

@Test
public void testFailWhileRestarting() throws Exception {
	try (SlotPool slotPool = createSlotPoolImpl()) {
		TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
		final ExecutionGraph executionGraph = TestingExecutionGraphBuilder.newBuilder()
			.setRestartStrategy(new InfiniteDelayRestartStrategy())
			.setTaskManagerLocation(taskManagerLocation)
			.buildAndScheduleForExecution(slotPool);

		// Release the TaskManager and wait for the job to restart
		slotPool.releaseTaskManager(taskManagerLocation.getResourceID(), new Exception("Test Exception"));

		assertEquals(JobStatus.RESTARTING, executionGraph.getState());

		// If we fail when being in RESTARTING, then we should try to restart again
		final long globalModVersion = executionGraph.getGlobalModVersion();
		final Exception testException = new Exception("Test exception");
		executionGraph.failGlobal(testException);

		assertNotEquals(globalModVersion, executionGraph.getGlobalModVersion());
		assertEquals(JobStatus.RESTARTING, executionGraph.getState());
		assertEquals(testException, executionGraph.getFailureCause()); // we should have updated the failure cause

		// but it should fail when sending a SuppressRestartsException
		executionGraph.failGlobal(new SuppressRestartsException(new Exception("Suppress restart exception")));

		assertEquals(JobStatus.FAILED, executionGraph.getState());

		// The restart has been aborted
		executionGraph.restart(executionGraph.getGlobalModVersion());

		assertEquals(JobStatus.FAILED, executionGraph.getState());
	}
}
 

@Nonnull
private static TestingLogicalSlot createTestingSlot(@Nullable CompletableFuture<?> releaseFuture) {
	return new TestingLogicalSlot(
		new LocalTaskManagerLocation(),
		new SimpleAckingTaskManagerGateway(),
		0,
		new AllocationID(),
		new SlotRequestId(),
		new SlotSharingGroupId(),
		releaseFuture);
}
 

/**
 * Tests that the root slot are moved from unresolved to resolved once the
 * slot context future is successfully completed
 */
@Test
public void testRootSlotTransition() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	assertTrue(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertFalse(slotSharingManager.getResolvedRootSlots().contains(rootSlot));

	// now complete the slotContextFuture
	slotContextFuture.complete(
		new SimpleSlotContext(
			new AllocationID(),
			new LocalTaskManagerLocation(),
			0,
			new SimpleAckingTaskManagerGateway()));

	assertFalse(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertTrue(slotSharingManager.getResolvedRootSlots().contains(rootSlot));
}
 

@Before
public void setUp() throws Exception {
	this.jobId = new JobID();

	taskManagerLocation = new LocalTaskManagerLocation();
	taskManagerGateway = new SimpleAckingTaskManagerGateway();
	resourceManagerGateway = new TestingResourceManagerGateway();
}
 

@Test
public void testSingleQueuedSharedSlotScheduling() throws Exception {
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture = slotProvider.allocateSlot(
		new ScheduledUnit(
			new JobVertexID(),
			slotSharingGroupId,
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

	boolean booleanCompletableFuture = slotPool.offerSlot(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		new SlotOffer(
			allocationId,
			0,
			ResourceProfile.UNKNOWN));

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

	assertEquals(slotSharingGroupId, logicalSlot.getSlotSharingGroupId());
}
 

/**
 * SlotPool#failAllocation should not fail with a {@link java.util.ConcurrentModificationException}
 * if there is a concurrent scheduling operation. See FLINK-13421.
 */
@Test
public void slotPoolExecutionGraph_ConcurrentSchedulingAndAllocationFailure_ShouldNotFailWithConcurrentModificationException() throws Exception {
	final SlotSharingGroup group = new SlotSharingGroup();
	final JobVertex vertex1 = createNoOpVertex("vertex1", 1);
	vertex1.setSlotSharingGroup(group);
	final JobVertex vertex2 = createNoOpVertex("vertex2", 3);
	vertex2.setSlotSharingGroup(group);
	final JobVertex vertex3 = createNoOpVertex("vertex3", 1);
	vertex3.setSlotSharingGroup(group);
	vertex3.connectNewDataSetAsInput(vertex2, DistributionPattern.ALL_TO_ALL, ResultPartitionType.PIPELINED);

	try (SlotPool slotPool = createSlotPoolImpl()) {
		final SlotProvider slots = createSchedulerWithSlots(slotPool, new LocalTaskManagerLocation(), 2);

		final AllocationID allocationId = slotPool.getAvailableSlotsInformation().iterator().next().getAllocationId();

		final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Test Job", vertex1, vertex2, vertex3);
		jobGraph.setScheduleMode(ScheduleMode.EAGER);
		final ExecutionGraph eg = TestingExecutionGraphBuilder
			.newBuilder()
			.setJobGraph(jobGraph)
			.setSlotProvider(slots)
			.setAllocationTimeout(Time.minutes(60))
			.build();

		startAndScheduleExecutionGraph(eg);

		slotPool.failAllocation(
			allocationId,
			new Exception("test exception"));

		eg.waitUntilTerminal();
	}
}
 

@Test
public void testSingleQueuedSharedSlotScheduling() throws Exception {
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture = slotProvider.allocateSlot(
		new ScheduledUnit(
			new JobVertexID(),
			slotSharingGroupId,
			null),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

	boolean booleanCompletableFuture = slotPool.offerSlot(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		new SlotOffer(
			allocationId,
			0,
			ResourceProfile.ANY));

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

	assertEquals(slotSharingGroupId, logicalSlot.getSlotSharingGroupId());
}
 

@Before
public void setUp() throws Exception {
	this.jobId = new JobID();

	taskManagerLocation = new LocalTaskManagerLocation();
	taskManagerGateway = new SimpleAckingTaskManagerGateway();
	resourceManagerGateway = new TestingResourceManagerGateway();
}
 

private SimpleSlotContext createSimpleSlotContext() {
	return new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());
}
 

private static ArchivedExecutionJobVertex createArchivedExecutionJobVertex(JobVertexID jobVertexID) {
	final StringifiedAccumulatorResult[] emptyAccumulators = new StringifiedAccumulatorResult[0];
	final long[] timestamps = new long[ExecutionState.values().length];
	final ExecutionState expectedState = ExecutionState.RUNNING;

	final LocalTaskManagerLocation assignedResourceLocation = new LocalTaskManagerLocation();
	final AllocationID allocationID = new AllocationID();

	final int subtaskIndex = 1;
	final int attempt = 2;
	return new ArchivedExecutionJobVertex(
		new ArchivedExecutionVertex[]{
			new ArchivedExecutionVertex(
				subtaskIndex,
				"test task",
				new ArchivedExecution(
					new StringifiedAccumulatorResult[0],
					null,
					new ExecutionAttemptID(),
					attempt,
					expectedState,
					"error",
					assignedResourceLocation,
					allocationID,
					subtaskIndex,
					timestamps),
				new EvictingBoundedList<>(0)
			)
		},
		jobVertexID,
		jobVertexID.toString(),
		1,
		1,
		ResourceProfile.UNKNOWN,
		emptyAccumulators);
}
 

@Test
public void testCalculationOfTaskExecutorUtilization() throws Exception {
	try (final SlotPoolImpl slotPool = createSlotPoolImpl()) {
		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);

		final TaskManagerLocation firstTaskManagerLocation = new LocalTaskManagerLocation();
		final TaskManagerLocation secondTaskManagerLocation = new LocalTaskManagerLocation();

		final List<AllocationID> firstTaskManagersSlots = registerAndOfferSlots(firstTaskManagerLocation, slotPool, 4);
		final List<AllocationID> secondTaskManagersSlots = registerAndOfferSlots(secondTaskManagerLocation, slotPool, 4);

		slotPool.allocateAvailableSlot(new SlotRequestId(), firstTaskManagersSlots.get(0));
		slotPool.allocateAvailableSlot(new SlotRequestId(), firstTaskManagersSlots.get(1));
		slotPool.allocateAvailableSlot(new SlotRequestId(), secondTaskManagersSlots.get(3));

		final Collection<SlotInfoWithUtilization> availableSlotsInformation = slotPool.getAvailableSlotsInformation();

		final Map<TaskManagerLocation, Double> utilizationPerTaskExecutor = ImmutableMap.of(
			firstTaskManagerLocation, 2.0 / 4,
			secondTaskManagerLocation, 1.0 / 4);

		for (SlotInfoWithUtilization slotInfoWithUtilization : availableSlotsInformation) {
			final double expectedTaskExecutorUtilization = utilizationPerTaskExecutor.get(slotInfoWithUtilization.getTaskManagerLocation());
			assertThat(slotInfoWithUtilization.getTaskExecutorUtilization(), is(closeTo(expectedTaskExecutorUtilization, 0.1)));
		}
	}
}
 

private void runSlotRequestCompletionTest(
	Supplier<SlotPoolImpl> slotPoolSupplier,
	Consumer<SlotPoolImpl> actionAfterSlotRequest) throws Exception {
	try (final SlotPoolImpl slotPool = slotPoolSupplier.get()) {

		final List<SlotRequestId> slotRequestIds = IntStream.range(0, 10)
			.mapToObj(ignored -> new SlotRequestId())
			.collect(Collectors.toList());

		final List<CompletableFuture<PhysicalSlot>> slotRequests = slotRequestIds
			.stream()
			.map(slotRequestId -> slotPool.requestNewAllocatedSlot(slotRequestId, ResourceProfile.UNKNOWN, TIMEOUT))
			.collect(Collectors.toList());

		actionAfterSlotRequest.accept(slotPool);

		final LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
		slotPool.registerTaskManager(taskManagerLocation.getResourceID());

		final SlotOffer slotOffer = new SlotOffer(new AllocationID(), 0, ResourceProfile.UNKNOWN);
		final Collection<SlotOffer> acceptedSlots = slotPool.offerSlots(taskManagerLocation, new SimpleAckingTaskManagerGateway(), Collections.singleton(slotOffer));

		assertThat(acceptedSlots, containsInAnyOrder(slotOffer));

		final FlinkException testingReleaseException = new FlinkException("Testing release exception");

		// check that the slot requests get completed in sequential order
		for (int i = 0; i < slotRequestIds.size(); i++) {
			final CompletableFuture<PhysicalSlot> slotRequestFuture = slotRequests.get(i);
			slotRequestFuture.get();
			slotPool.releaseSlot(slotRequestIds.get(i), testingReleaseException);
		}
	}
}
 

/**
 * 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()));
}
 

public static ResourceID offerSlots(
		SlotPoolImpl slotPool,
		ComponentMainThreadExecutor mainThreadExecutor,
		List<ResourceProfile> resourceProfiles,
		TaskManagerGateway taskManagerGateway) {
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	CompletableFuture.runAsync(
		() -> {
			slotPool.registerTaskManager(taskManagerLocation.getResourceID());

			final Collection<SlotOffer> slotOffers = IntStream
				.range(0, resourceProfiles.size())
				.mapToObj(i -> new SlotOffer(new AllocationID(), i, resourceProfiles.get(i)))
				.collect(Collectors.toList());

			final Collection<SlotOffer> acceptedOffers = slotPool.offerSlots(
				taskManagerLocation,
				taskManagerGateway,
				slotOffers);

			assertThat(acceptedOffers, is(slotOffers));
		},
		mainThreadExecutor
	).join();

	return taskManagerLocation.getResourceID();
}
 

@Before
public void setUp() throws Exception {
	this.jobId = new JobID();

	taskManagerLocation = new LocalTaskManagerLocation();
	taskManagerGateway = new SimpleAckingTaskManagerGateway();
	resourceManagerGateway = new TestingResourceManagerGateway();
}
 

@Test
public void testGetResolvedSlotWithResourceConfigured() {
	ResourceProfile rp1 = ResourceProfile.fromResources(1.0, 100);
	ResourceProfile rp2 = ResourceProfile.fromResources(2.0, 200);
	ResourceProfile allocatedSlotRp = ResourceProfile.fromResources(5.0, 500);

	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			CompletableFuture.completedFuture(
					new SimpleSlotContext(
							new AllocationID(),
							new LocalTaskManagerLocation(),
							0,
							new SimpleAckingTaskManagerGateway(),
							allocatedSlotRp)),
			new SlotRequestId());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	Collection<SlotSelectionStrategy.SlotInfoAndResources> resolvedRoots =
		slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1), resolvedRoots.iterator().next().getRemainingResources());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp2,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	resolvedRoots = slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1).subtract(rp2), resolvedRoots.iterator().next().getRemainingResources());
}
 

/**
 * Tests that incomplete futures returned by {@link ShuffleMaster#registerPartitionWithProducer} are rejected.
 */
@Test
public void testIncompletePartitionRegistrationFutureIsRejected() throws Exception {
	final ShuffleMaster<ShuffleDescriptor> shuffleMaster = new TestingShuffleMaster();
	final JobGraph jobGraph = new JobGraph("job graph");
	final JobVertex source = new JobVertex("source");
	final JobVertex target = new JobVertex("target");

	source.setInvokableClass(AbstractInvokable.class);
	target.setInvokableClass(AbstractInvokable.class);
	target.connectNewDataSetAsInput(source, POINTWISE, PIPELINED);
	jobGraph.addVertex(source);
	jobGraph.addVertex(target);
	ExecutionGraph executionGraph = TestingExecutionGraphBuilder
		.newBuilder()
		.setJobGraph(jobGraph)
		.setShuffleMaster(shuffleMaster)
		.build();

	final ExecutionVertex sourceVertex = executionGraph.getAllVertices().get(source.getID()).getTaskVertices()[0];

	boolean incompletePartitionRegistrationRejected = false;
	try {
		Execution.registerProducedPartitions(sourceVertex, new LocalTaskManagerLocation(), new ExecutionAttemptID(), false);
	} catch (IllegalStateException e) {
		incompletePartitionRegistrationRejected = true;
	}

	assertTrue(incompletePartitionRegistrationRejected);
}
 

public TestingLogicalSlot(TaskManagerGateway taskManagerGateway) {
	this(
		new LocalTaskManagerLocation(),
		taskManagerGateway,
		0,
		new AllocationID(),
		new SlotRequestId(),
		new SlotSharingGroupId(),
		null);
}
 

@Test
public void testTaskExecutorUtilizationCalculation() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();
	final TaskManagerLocation firstTaskExecutorLocation = new LocalTaskManagerLocation();
	final TaskManagerLocation secondTaskExecutorLocation = new LocalTaskManagerLocation();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotSharingManager.MultiTaskSlot firstRootSlot = createRootSlot(firstTaskExecutorLocation, slotSharingManager);
	createRootSlot(firstTaskExecutorLocation, slotSharingManager);
	createRootSlot(secondTaskExecutorLocation, slotSharingManager);

	final AbstractID groupId = new AbstractID();

	firstRootSlot.allocateSingleTaskSlot(new SlotRequestId(), ResourceProfile.UNKNOWN, groupId, Locality.UNCONSTRAINED);

	final Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfoAndResources = slotSharingManager.listResolvedRootSlotInfo(groupId);

	assertThat(slotInfoAndResources, hasSize(2));

	final Map<TaskManagerLocation, Double> utilizationPerTaskExecutor = slotInfoAndResources.stream()
		.collect(
			Collectors.toMap(
				slot -> slot.getSlotInfo().getTaskManagerLocation(),
				SlotSelectionStrategy.SlotInfoAndResources::getTaskExecutorUtilization));

	assertThat(utilizationPerTaskExecutor.get(firstTaskExecutorLocation), is(closeTo(1.0 / 2, 0.1)));
	assertThat(utilizationPerTaskExecutor.get(secondTaskExecutorLocation), is(closeTo(0, 0.1)));
}
 

@Test
public void testSingleQueuedSharedSlotScheduling() throws Exception {
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final SlotPoolImpl slotPool = slotPoolResource.getSlotPool();
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture = slotProvider.allocateSlot(
		new ScheduledUnit(
			new JobVertexID(),
			slotSharingGroupId,
			null),
		true,
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	assertFalse(logicalSlotFuture.isDone());

	final AllocationID allocationId = allocationIdFuture.get();

	boolean booleanCompletableFuture = slotPool.offerSlot(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		new SlotOffer(
			allocationId,
			0,
			ResourceProfile.UNKNOWN));

	assertTrue(booleanCompletableFuture);

	final LogicalSlot logicalSlot = logicalSlotFuture.get();

	assertEquals(slotSharingGroupId, logicalSlot.getSlotSharingGroupId());
}
 

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;
	});
}
 

/**
 * Tests that we cannot retrieve a slot when it's releasing children.
 */
@Test
public void testResolvedSlotInReleasingIsNotAvailable() throws Exception {
	final SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	final SlotSharingManager.MultiTaskSlot rootSlot = createRootSlot(new LocalTaskManagerLocation(), slotSharingManager);

	final AbstractID groupId1 = new AbstractID();
	final SlotSharingManager.SingleTaskSlot singleTaskSlot = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId1,
		Locality.UNCONSTRAINED);

	final AtomicBoolean verified = new AtomicBoolean(false);

	final AbstractID groupId2 = new AbstractID();
	// register a verification in MultiTaskSlot's children releasing loop
	singleTaskSlot.getLogicalSlotFuture().get().tryAssignPayload(new LogicalSlot.Payload() {
		@Override
		public void fail(Throwable cause) {
			assertEquals(0, slotSharingManager.listResolvedRootSlotInfo(groupId2).size());

			verified.set(true);
		}

		@Override
		public CompletableFuture<?> getTerminalStateFuture() {
			return null;
		}
	});

	assertEquals(1, slotSharingManager.listResolvedRootSlotInfo(groupId2).size());

	rootSlot.release(new Exception("test exception"));

	// ensure the verification in Payload#fail is passed
	assertTrue(verified.get());
}