Java Code Examples for org.apache.flink.runtime.clusterframework.types.ResourceProfile#fromResources()

@Test
public void testUnfulfillableRequestsFailWhenOn() {
	// setup
	final ResourceProfile availableProfile = ResourceProfile.fromResources(2.0, 100);
	final ResourceProfile unfulfillableProfile = ResourceProfile.fromResources(2.0, 200);

	final List<Tuple3<JobID, AllocationID, Exception>> notifiedAllocationFailures = new ArrayList<>();
	final SlotManager slotManager = createSlotManagerNotStartingNewTMs(notifiedAllocationFailures);
	registerFreeSlot(slotManager, availableProfile);

	// test
	try {
		slotManager.registerSlotRequest(slotRequest(unfulfillableProfile));
		fail("this should cause an exception");
	} catch (ResourceManagerException exception) {
		assertTrue(ExceptionUtils.findThrowable(exception, UnfulfillableSlotRequestException.class).isPresent());
	}

	// assert
	assertEquals(0, notifiedAllocationFailures.size());
	assertEquals(0, slotManager.getNumberPendingSlotRequests());
}
 

/**
 * Tests that the slot manager ignores slot reports of unknown origin (not registered
 * task managers).
 */
@Test
public void testReceivingUnknownSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotStatus unknownSlotStatus = new SlotStatus(unknownSlotId, unknownResourceProfile);
	final SlotReport unknownSlotReport = new SlotReport(unknownSlotStatus);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		// check that we don't have any slots registered
		assertTrue(0 == slotManager.getNumberRegisteredSlots());

		// this should not update anything since the instance id is not known to the slot manager
		assertFalse(slotManager.reportSlotStatus(unknownInstanceID, unknownSlotReport));

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
	}
}
 

@Test
public void testTurnOnKeepsPendingFulfillableRequests() throws Exception {
	// setup
	final ResourceProfile availableProfile = ResourceProfile.fromResources(2.0, 100);
	final ResourceProfile fulfillableProfile = ResourceProfile.fromResources(1.0, 100);

	final SlotManager slotManager = createSlotManagerNotStartingNewTMs();
	slotManager.setFailUnfulfillableRequest(false);
	registerFreeSlot(slotManager, availableProfile);

	slotManager.registerSlotRequest(slotRequest(fulfillableProfile));
	slotManager.registerSlotRequest(slotRequest(fulfillableProfile));

	// test
	slotManager.setFailUnfulfillableRequest(true);

	// assert
	assertEquals(1, slotManager.getNumberPendingSlotRequests());
}
 

@Test
public void testTurnOnKeepsRequestsWithStartingTMs() throws Exception {
	// setup
	final ResourceProfile availableProfile = ResourceProfile.fromResources(2.0, 100);
	final ResourceProfile newTmProfile = ResourceProfile.fromResources(2.0, 200);

	final SlotManager slotManager = createSlotManagerStartingNewTMs();
	slotManager.setFailUnfulfillableRequest(false);
	registerFreeSlot(slotManager, availableProfile);

	// test
	slotManager.registerSlotRequest(slotRequest(newTmProfile));
	slotManager.setFailUnfulfillableRequest(true);

	// assert
	assertEquals(1, slotManager.getNumberPendingSlotRequests());
}
 

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

	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();

	SlotSharingManager.MultiTaskSlot unresolvedRootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			slotContextFuture,
			new SlotRequestId());

	SlotSharingManager.MultiTaskSlot multiTaskSlot =
			unresolvedRootSlot.allocateMultiTaskSlot(new SlotRequestId(), new SlotSharingGroupId());

	SlotSharingManager.SingleTaskSlot firstChild = multiTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(rp1), is(true));
	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(rp2), is(true));
	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(ResourceProfile.UNKNOWN), is(true));

	slotContextFuture.complete(new AllocatedSlot(
			new AllocationID(),
			new TaskManagerLocation(new ResourceID("tm-X"), InetAddress.getLoopbackAddress(), 46),
			0,
			allocatedSlotRp,
			mock(TaskManagerGateway.class)));

	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(rp1), is(true));
	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(rp2), is(false));
	assertThat(multiTaskSlot.mayHaveEnoughResourcesToFulfill(ResourceProfile.UNKNOWN), is(true));
}
 

/**
 * Tests that duplicate slot requests (requests with an already registered allocation id) are
 * also detected after a pending slot request has been fulfilled but not yet freed.
 */
@Test
public void testDuplicatePendingSlotRequestAfterSuccessfulAllocation() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final AtomicInteger allocateResourceCalls = new AtomicInteger(0);
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceConsumer(ignored -> allocateResourceCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = ResourceProfile.fromResources(1.0, 2);
	final ResourceProfile resourceProfile2 = ResourceProfile.fromResources(2.0, 1);
	final SlotRequest slotRequest1 = new SlotRequest(new JobID(), allocationId, resourceProfile1, "foobar");
	final SlotRequest slotRequest2 = new SlotRequest(new JobID(), allocationId, resourceProfile2, "barfoo");

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();

	final ResourceID resourceID = ResourceID.generate();

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile1);
	final SlotReport slotReport = new SlotReport(slotStatus);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);
		assertTrue(slotManager.registerSlotRequest(slotRequest1));

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());

		assertFalse(slotManager.registerSlotRequest(slotRequest2));
	}

	// check that we have only called the resource allocation only for the first slot request,
	// since the second request is a duplicate
	assertThat(allocateResourceCalls.get(), is(0));
}
 

/**
 * Tests that freeing a slot will correctly reset the slot and mark it as a free slot.
 */
@Test
public void testFreeSlot() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final JobID jobId = new JobID();

	ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
	final ResourceID resourceID = taskExecutorConnection.getResourceID();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(42.0, 1337);

	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile, jobId, allocationId);
	final SlotReport slotReport = new SlotReport(slotStatus);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {

		slotManager.registerTaskManager(
			taskExecutorConnection,
			slotReport);

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());

		// this should be ignored since the allocation id does not match
		slotManager.freeSlot(slotId, new AllocationID());

		assertTrue(slot.getState() == TaskManagerSlot.State.ALLOCATED);
		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());

		slotManager.freeSlot(slotId, allocationId);

		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
		assertNull(slot.getAllocationId());
	}
}
 

/**
 * Checks that un-registering a pending slot request will cancel it, removing it from all
 * assigned task manager slots and then remove it from the slot manager.
 */
@Test
public void testUnregisterPendingSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> new CompletableFuture<>())
		.createTestingTaskExecutorGateway();

	final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final SlotRequest slotRequest = new SlotRequest(new JobID(), allocationId, resourceProfile, "foobar");

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		slotManager.registerSlotRequest(slotRequest);

		assertNotNull(slotManager.getSlotRequest(allocationId));

		assertTrue(slot.getState() == TaskManagerSlot.State.PENDING);

		slotManager.unregisterSlotRequest(allocationId);

		assertNull(slotManager.getSlotRequest(allocationId));

		slot = slotManager.getSlot(slotId);
		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
	}
}
 

/**
 * Tests that a different slot can fulfill a pending slot request. If the
 * pending slot request has a pending task manager slot assigned, it should
 * be freed.
 */
@Test
public void testRegistrationOfDifferentSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder().build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions, numberSlots)) {
		final JobID jobId = new JobID();
		final ResourceProfile requestedSlotProfile = ResourceProfile.fromResources(1.0, 1);

		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId, requestedSlotProfile)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));

		final int numberOfferedSlots = 1;
		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final ResourceProfile offeredSlotProfile = ResourceProfile.fromResources(2.0, 2);
		final SlotReport slotReport = createSlotReport(
			taskExecutorConnection.getResourceID(),
			numberOfferedSlots,
			offeredSlotProfile,
			SlotManagerImplTest::createEmptySlotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberOfferedSlots));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(0));
	}
}
 

@Test
public void testFulfillableRequestsKeepPendingWhenOn() throws Exception {
	// setup
	final ResourceProfile availableProfile = ResourceProfile.fromResources(2.0, 100);

	final SlotManager slotManager = createSlotManagerNotStartingNewTMs();
	registerFreeSlot(slotManager, availableProfile);

	// test
	slotManager.registerSlotRequest(slotRequest(availableProfile));
	slotManager.registerSlotRequest(slotRequest(availableProfile));

	// assert
	assertEquals(1, slotManager.getNumberPendingSlotRequests());
}
 

@Test
public void testStartingTmKeepsSlotPendingWhenOn() throws Exception {
	// setup
	final ResourceProfile availableProfile = ResourceProfile.fromResources(2.0, 100);
	final ResourceProfile newTmProfile = ResourceProfile.fromResources(2.0, 200);

	final SlotManager slotManager = createSlotManagerStartingNewTMs();
	registerFreeSlot(slotManager, availableProfile);

	// test
	slotManager.registerSlotRequest(slotRequest(newTmProfile));

	// assert
	assertEquals(1, slotManager.getNumberPendingSlotRequests());
}
 

/**
 * Tests that a slot request is retried if it times out on the task manager side.
 */
@Test
public void testTaskManagerSlotRequestTimeoutHandling() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(42.0, 1337);
	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");
	final CompletableFuture<Acknowledge> slotRequestFuture1 = new CompletableFuture<>();
	final CompletableFuture<Acknowledge> slotRequestFuture2 = new CompletableFuture<>();
	final Iterator<CompletableFuture<Acknowledge>> slotRequestFutureIterator = Arrays.asList(slotRequestFuture1, slotRequestFuture2).iterator();
	final ArrayBlockingQueue<SlotID> slotIds = new ArrayBlockingQueue<>(2);

	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(FunctionUtils.uncheckedFunction(
			requestSlotParameters -> {
				slotIds.put(requestSlotParameters.f0);
				return slotRequestFutureIterator.next();
			}))
		.createTestingTaskExecutorGateway();

	final ResourceID resourceId = ResourceID.generate();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceId, taskExecutorGateway);

	final SlotID slotId1 = new SlotID(resourceId, 0);
	final SlotID slotId2 = new SlotID(resourceId, 1);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, resourceProfile);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, resourceProfile);
	final SlotReport slotReport = new SlotReport(Arrays.asList(slotStatus1, slotStatus2));

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {

		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		slotManager.registerSlotRequest(slotRequest);

		final SlotID firstSlotId = slotIds.take();
		assertThat(slotIds, is(empty()));

		TaskManagerSlot failedSlot = slotManager.getSlot(firstSlotId);

		// let the first attempt fail --> this should trigger a second attempt
		slotRequestFuture1.completeExceptionally(new SlotAllocationException("Test exception."));

		// the second attempt succeeds
		slotRequestFuture2.complete(Acknowledge.get());

		final SlotID secondSlotId = slotIds.take();
		assertThat(slotIds, is(empty()));

		TaskManagerSlot slot = slotManager.getSlot(secondSlotId);

		assertTrue(slot.getState() == TaskManagerSlot.State.ALLOCATED);
		assertEquals(allocationId, slot.getAllocationId());

		if (!failedSlot.getSlotId().equals(slot.getSlotId())) {
			assertTrue(failedSlot.getState() == TaskManagerSlot.State.FREE);
		}
	}
}
 

/**
 * Tests that a slot request which can be fulfilled will trigger a slot allocation.
 */
@Test
public void testSlotRequestWithFreeSlot() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final JobID jobId = new JobID();
	final SlotID slotId = new SlotID(resourceID, 0);
	final String targetAddress = "localhost";
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(42.0, 1337);
	final SlotRequest slotRequest = new SlotRequest(
		jobId,
		allocationId,
		resourceProfile,
		targetAddress);

	ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		final CompletableFuture<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestFuture = new CompletableFuture<>();
		// accept an incoming slot request
		final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setRequestSlotFunction(tuple6 -> {
				requestFuture.complete(Tuple6.of(tuple6.f0, tuple6.f1, tuple6.f2, tuple6.f3, tuple6.f4, tuple6.f5));
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.createTestingTaskExecutorGateway();

		final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

		final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
		final SlotReport slotReport = new SlotReport(slotStatus);

		slotManager.registerTaskManager(
			taskExecutorConnection,
			slotReport);

		assertTrue("The slot request should be accepted", slotManager.registerSlotRequest(slotRequest));

		assertThat(requestFuture.get(), is(equalTo(Tuple6.of(slotId, jobId, allocationId, resourceProfile, targetAddress, resourceManagerId))));

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());
	}
}
 

/**
 * 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 taskResourceProfile = ResourceProfile.fromResources(0.5, 250);
	final ResourceProfile physicalSlotResourceProfile = ResourceProfile.fromResources(1.0, 300);
	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final TestingStateLocationRetriever stateLocationRetriever = new TestingStateLocationRetriever();
	stateLocationRetriever.setStateLocation(executionVertexId, taskManagerLocation);

	final ExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator(
		stateLocationRetriever,
		new TestingInputsLocationsRetriever.Builder().build());

	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
		new ExecutionVertexSchedulingRequirements.Builder()
			.withExecutionVertexId(executionVertexId)
			.withPreviousAllocationId(allocationId)
			.withSlotSharingGroupId(sharingGroupId)
			.withPhysicalSlotResourceProfile(physicalSlotResourceProfile)
			.withTaskResourceProfile(taskResourceProfile)
			.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(taskResourceProfile, expectedSlotProfile.getTaskResourceProfile());
	assertEquals(physicalSlotResourceProfile, expectedSlotProfile.getPhysicalSlotResourceProfile());
	assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocation));
}
 

/**
 * Tests whether
 * 1) SlotManager accepts a slot request.
 * 2) SlotRequest leads to a container allocation.
 * 3) Slot becomes available and TaskExecutor gets a SlotRequest.
 */
@Test
public void testSlotsUnavailableRequest() throws Exception {
	final JobID jobID = new JobID();

	final ResourceManagerId rmLeaderID = ResourceManagerId.generate();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setDefaultWorkerResourceSpec(new WorkerResourceSpec.Builder()
			.setCpuCores(1.0)
			.setTaskHeapMemoryMB(100)
			.build())
		.setScheduledExecutor(scheduledExecutor)
		.build()) {

		final CompletableFuture<WorkerResourceSpec> workerRequestFuture = new CompletableFuture<>();
		ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
			.setAllocateResourceConsumer(workerRequestFuture::complete)
			.build();

		slotManager.start(rmLeaderID, Executors.directExecutor(), resourceManagerActions);

		final AllocationID allocationID = new AllocationID();
		final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 100);
		final String targetAddress = "foobar";

		SlotRequest slotRequest = new SlotRequest(jobID, allocationID, resourceProfile, targetAddress);

		slotManager.registerSlotRequest(slotRequest);

		workerRequestFuture.get();

		// slot becomes available
		final CompletableFuture<Tuple3<SlotID, JobID, AllocationID>> requestFuture = new CompletableFuture<>();
		TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setRequestSlotFunction(tuple5 -> {
				requestFuture.complete(Tuple3.of(tuple5.f0, tuple5.f1, tuple5.f2));
				return new CompletableFuture<>();
			})
			.createTestingTaskExecutorGateway();

		final ResourceID resourceID = ResourceID.generate();
		final SlotID slotID = new SlotID(resourceID, 0);

		final SlotStatus slotStatus =
			new SlotStatus(slotID, resourceProfile);
		final SlotReport slotReport =
			new SlotReport(Collections.singletonList(slotStatus));
		// register slot at SlotManager
		slotManager.registerTaskManager(new TaskExecutorConnection(resourceID, taskExecutorGateway), slotReport);

		// 4) Slot becomes available and TaskExecutor gets a SlotRequest
		assertThat(requestFuture.get(), is(equalTo(Tuple3.of(slotID, jobID, allocationID))));
	}
}
 

/**
 * Tests that the correct partition/slot report is sent as part of the heartbeat response.
 */
@Test
public void testHeartbeatReporting() throws Exception {
	final String rmAddress = "rm";
	final UUID rmLeaderId = UUID.randomUUID();

	// register the mock resource manager gateway
	final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
	final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
	final ResourceID rmResourceId = rmGateway.getOwnResourceId();
	final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmResourceId,
			new ClusterInformation("localhost", 1234)));

	rmGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
		taskExecutorRegistrationFuture.complete(taskExecutorRegistration.getResourceId());
		return registrationResponse;
	});

	final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
	rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<TaskExecutorHeartbeatPayload> heartbeatPayloadCompletableFuture = new CompletableFuture<>();
	rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, heartbeatPayload) -> heartbeatPayloadCompletableFuture.complete(heartbeatPayload));

	rpc.registerGateway(rmAddress, rmGateway);

	final SlotID slotId = new SlotID(unresolvedTaskManagerLocation.getResourceID(), 0);
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotReport slotReport1 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile));
	final SlotReport slotReport2 = new SlotReport(
		new SlotStatus(
			slotId,
			resourceProfile,
			new JobID(),
			new AllocationID()));

	final Queue<SlotReport> reports = new ArrayDeque<>(Arrays.asList(slotReport1, slotReport2));
	final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable
		.<Task>newBuilder()
		.createSlotReportSupplier(reports::poll)
		.closeAsyncReturns(CompletableFuture.completedFuture(null))
		.build();

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutorPartitionTracker partitionTracker = createPartitionTrackerWithFixedPartitionReport(taskManagerServices.getShuffleEnvironment());

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices, HEARTBEAT_SERVICES, partitionTracker);

	try {
		taskManager.start();

		// define a leader and see that a registration happens
		resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);

		// register resource manager success will trigger monitoring heartbeat target between tm and rm
		assertThat(taskExecutorRegistrationFuture.get(), equalTo(unresolvedTaskManagerLocation.getResourceID()));
		assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));

		TaskExecutorGateway taskExecutorGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// trigger the heartbeat asynchronously
		taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);

		// wait for heartbeat response
		SlotReport actualSlotReport = heartbeatPayloadCompletableFuture.get().getSlotReport();

		// the new slot report should be reported
		assertEquals(slotReport2, actualSlotReport);

		ClusterPartitionReport actualClusterPartitionReport = heartbeatPayloadCompletableFuture.get().getClusterPartitionReport();
		assertEquals(partitionTracker.createClusterPartitionReport(), actualClusterPartitionReport);
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that when matching from the allocated slot, the remaining resources of the slot
 * will be used instead of the total resource.
 */
@Test
public void testSlotSharingRespectsRemainingResource() throws Exception {
	final ResourceProfile allocatedSlotRp = ResourceProfile.fromResources(3.0, 300);
	final ResourceProfile largeRequestResource = ResourceProfile.fromResources(2.0, 200);
	final ResourceProfile smallRequestResource = ResourceProfile.fromResources(1.0, 100);

	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);
	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

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

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();
	final JobVertexID jobVertexId1 = new JobVertexID();
	final JobVertexID jobVertexId2 = new JobVertexID();
	final JobVertexID jobVertexId3 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture1 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId1,
					slotSharingGroupId,
					null),
			SlotProfile.noLocality(largeRequestResource),
			TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();

	// This should fulfill the first request.
	boolean offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId1,
					0,
					allocatedSlotRp));

	assertTrue(offerFuture);
	assertTrue(logicalSlotFuture1.isDone());
	assertEquals(allocationId1, logicalSlotFuture1.get().getAllocationId());

	// The second request should not share the same slot with the first request since it is large.
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId2,
					slotSharingGroupId,
					null),
			SlotProfile.noLocality(largeRequestResource),
			TestingUtils.infiniteTime());
	assertFalse(logicalSlotFuture2.isDone());

	// The third request should be able to share the same slot with the first request since it is small.
	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId3,
					slotSharingGroupId,
					null),
			SlotProfile.noLocality(smallRequestResource),
			TestingUtils.infiniteTime());
	assertTrue(logicalSlotFuture3.isDone());
	assertEquals(allocationId1, logicalSlotFuture1.get().getAllocationId());

	// The second request should be finally fulfilled by a new slot.
	final AllocationID allocationId2 = allocationIds.take();
	// This should fulfill the first two requests.
	offerFuture = slotPool.offerSlot(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			new SlotOffer(
					allocationId2,
					0,
					allocatedSlotRp));

	assertTrue(offerFuture);
	assertTrue(logicalSlotFuture2.isDone());
	assertEquals(allocationId2, logicalSlotFuture2.get().getAllocationId());
}
 

/**
 * Tests that an already registered allocation id can be reused after the initial slot request
 * has been freed.
 */
@Test
public void testAcceptingDuplicateSlotRequestAfterAllocationRelease() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final AtomicInteger allocateResourceCalls = new AtomicInteger(0);
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceConsumer(ignored -> allocateResourceCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = ResourceProfile.fromResources(1.0, 2);
	final ResourceProfile resourceProfile2 = ResourceProfile.fromResources(2.0, 1);
	final SlotRequest slotRequest1 = new SlotRequest(new JobID(), allocationId, resourceProfile1, "foobar");
	final SlotRequest slotRequest2 = new SlotRequest(new JobID(), allocationId, resourceProfile2, "barfoo");

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();

	final ResourceID resourceID = ResourceID.generate();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final SlotStatus slotStatus = new SlotStatus(slotId, ResourceProfile.fromResources(2.0, 2));
	final SlotReport slotReport = new SlotReport(slotStatus);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);
		assertTrue(slotManager.registerSlotRequest(slotRequest1));

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());

		slotManager.freeSlot(slotId, allocationId);

		// check that the slot has been freed
		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
		assertNull(slot.getAllocationId());

		assertTrue(slotManager.registerSlotRequest(slotRequest2));

		assertEquals("The slot has not been allocated to the expected allocation id.", allocationId, slot.getAllocationId());
	}

	// check that we have only called the resource allocation only for the first slot request,
	// since the second request is a duplicate
	assertThat(allocateResourceCalls.get(), is(0));
}
 

private SlotSharingResourceTestContext createResourceTestContext(ResourceProfile allocatedResourceProfile) {
	ResourceProfile coLocationTaskRp = ResourceProfile.fromResources(2.0, 200);
	ResourceProfile thirdChildRp = ResourceProfile.fromResources(3.0, 300);
	ResourceProfile forthChildRp = ResourceProfile.fromResources(9.0, 900);

	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();

	SlotSharingManager.MultiTaskSlot unresolvedRootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			slotContextFuture,
			new SlotRequestId());

	SlotSharingManager.MultiTaskSlot coLocationTaskSlot = unresolvedRootSlot.allocateMultiTaskSlot(
			new SlotRequestId(), new SlotSharingGroupId());

	SlotSharingManager.SingleTaskSlot firstCoLocatedChild = coLocationTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			coLocationTaskRp,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	SlotSharingManager.SingleTaskSlot secondCoLocatedChild = coLocationTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			coLocationTaskRp,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	SlotSharingManager.SingleTaskSlot thirdChild = unresolvedRootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			thirdChildRp,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	SlotSharingManager.SingleTaskSlot forthChild = unresolvedRootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			forthChildRp,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	slotContextFuture.complete(new AllocatedSlot(
			new AllocationID(),
			new TaskManagerLocation(new ResourceID("tm-X"), InetAddress.getLoopbackAddress(), 46),
			0,
			allocatedResourceProfile,
			mock(TaskManagerGateway.class)));

	return new SlotSharingResourceTestContext(
			slotSharingManager,
			coLocationTaskSlot,
			Arrays.asList(firstCoLocatedChild, secondCoLocatedChild, thirdChild, forthChild));
}
 

/**
 * Tests that slot requests time out after the specified request timeout. If a slot request
 * times out, then the request is cancelled, removed from the slot manager and the resource
 * manager is notified about the failed allocation.
 */
@Test
public void testSlotRequestTimeout() throws Exception {
	final long allocationTimeout = 50L;

	final CompletableFuture<Tuple2<JobID, AllocationID>> failedAllocationFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setNotifyAllocationFailureConsumer(tuple3 -> failedAllocationFuture.complete(Tuple2.of(tuple3.f0, tuple3.f1)))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();

	final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");

	final Executor mainThreadExecutor = TestingUtils.defaultExecutor();

	try (SlotManager slotManager = createSlotManagerBuilder()
		.setSlotRequestTimeout(Time.milliseconds(allocationTimeout))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		final AtomicReference<Exception> atomicException = new AtomicReference<>(null);

		mainThreadExecutor.execute(() -> {
			try {
				assertTrue(slotManager.registerSlotRequest(slotRequest));
			} catch (Exception e) {
				atomicException.compareAndSet(null, e);
			}
		});

		assertThat(failedAllocationFuture.get(), is(equalTo(Tuple2.of(jobId, allocationId))));

		if (atomicException.get() != null) {
			throw atomicException.get();
		}
	}
}