Java Code Examples for org.apache.flink.runtime.resourcemanager.ResourceManagerId#generate()

/**
 * Tests that if we have received a slot report with some allocated slots, then we don't accept
 * slot requests with allocated allocation ids.
 */
@Test
public void testDuplicatePendingSlotRequestAfterSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);

	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

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

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

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

		assertFalse(slotManager.registerSlotRequest(slotRequest));
	}
}
 

/**
 * Tests that we can register task manager and their slots at the slot manager.
 */
@Test
public void testTaskManagerRegistration() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().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 ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, resourceProfile);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, resourceProfile);
	final SlotReport slotReport = new SlotReport(Arrays.asList(slotStatus1, slotStatus2));

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

		assertTrue("The number registered slots does not equal the expected number.", 2 == slotManager.getNumberRegisteredSlots());

		assertNotNull(slotManager.getSlot(slotId1));
		assertNotNull(slotManager.getSlot(slotId2));
	}
}
 

/**
 * Tests that we can register task manager and their slots at the slot manager.
 */
@Test
public void testTaskManagerRegistration() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().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 ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	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);

		assertTrue("The number registered slots does not equal the expected number.", 2 == slotManager.getNumberRegisteredSlots());

		assertNotNull(slotManager.getSlot(slotId1));
		assertNotNull(slotManager.getSlot(slotId2));
	}
}
 

/**
 * 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 = mock(ResourceActions.class);

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = new ResourceProfile(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());
	}
}
 

/**
 * 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(resourceProfile -> allocateResourceCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = new ResourceProfile(1.0, 2);
	final ResourceProfile resourceProfile2 = new ResourceProfile(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 idle task managers time out after the configured timeout. A timed out task manager
 * will be removed from the slot manager and the resource manager will be notified about the
 * timeout, if it can be released.
 */
@Test
public void testTaskManagerTimeout() throws Exception {
	final long tmTimeout = 10L;

	final CompletableFuture<InstanceID> releaseFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releaseFuture.complete(instanceID))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final Executor mainThreadExecutor = TestingUtils.defaultExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(Time.milliseconds(tmTimeout))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		mainThreadExecutor.execute(() -> slotManager.registerTaskManager(taskManagerConnection, slotReport));

		assertThat(releaseFuture.get(), is(equalTo(taskManagerConnection.getInstanceID())));
	}
}
 

/**
 * Tests that the slot request fails if we cannot allocate more resources.
 */
@Test
public void testSlotRequestWithResourceAllocationFailure() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	final SlotRequest slotRequest = new SlotRequest(
		new JobID(),
		new AllocationID(),
		resourceProfile,
		"localhost");

	ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(value -> {
			throw new ResourceManagerException("Test exception");
		})
		.build();

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

		slotManager.registerSlotRequest(slotRequest);

		fail("The slot request should have failed with a ResourceManagerException.");

	} catch (ResourceManagerException e) {
		// expected exception
	}
}
 

/**
 * Tests that a task manager timeout does not remove the slots from the SlotManager.
 * A timeout should only trigger the {@link ResourceActions#releaseResource(InstanceID, Exception)}
 * callback. The receiver of the callback can then decide what to do with the TaskManager.
 *
 * <p>See FLINK-7793
 */
@Test
public void testTaskManagerTimeoutDoesNotRemoveSlots() throws Exception {
	final Time taskManagerTimeout = Time.milliseconds(10L);
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final ResourceActions resourceActions = mock(ResourceActions.class);
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

	when(taskExecutorGateway.canBeReleased()).thenReturn(CompletableFuture.completedFuture(true));

	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
	final SlotStatus slotStatus = new SlotStatus(
		new SlotID(resourceID, 0),
		new ResourceProfile(1.0, 1));
	final SlotReport initialSlotReport = new SlotReport(slotStatus);

	try (final SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.build()) {

		slotManager.start(resourceManagerId, Executors.directExecutor(), resourceActions);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		verify(resourceActions, timeout(taskManagerTimeout.toMilliseconds() * 20L).atLeast(1)).releaseResource(eq(taskExecutorConnection.getInstanceID()), any(Exception.class));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID());

		assertEquals(0, slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * 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 = mock(ResourceActions.class);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

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

	final ResourceProfile resourceProfile = new ResourceProfile(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 (SlotManager 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);
	}
}
 

public TestingResourceManagerGateway() {
	this(
		ResourceManagerId.generate(),
		ResourceID.generate(),
		"localhost/" + UUID.randomUUID(),
		"localhost");
}
 

/**
 * Tests that a task manager timeout does not remove the slots from the SlotManager.
 * A timeout should only trigger the {@link ResourceActions#releaseResource(InstanceID, Exception)}
 * callback. The receiver of the callback can then decide what to do with the TaskManager.
 *
 * <p>See FLINK-7793
 */
@Test
public void testTaskManagerTimeoutDoesNotRemoveSlots() throws Exception {
	final Time taskManagerTimeout = Time.milliseconds(10L);
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final ResourceActions resourceActions = mock(ResourceActions.class);
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

	when(taskExecutorGateway.canBeReleased()).thenReturn(CompletableFuture.completedFuture(true));

	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
	final SlotStatus slotStatus = new SlotStatus(
		new SlotID(resourceID, 0),
		new ResourceProfile(1.0, 1));
	final SlotReport initialSlotReport = new SlotReport(slotStatus);

	try (final SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.build()) {

		slotManager.start(resourceManagerId, Executors.directExecutor(), resourceActions);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		verify(resourceActions, timeout(taskManagerTimeout.toMilliseconds() * 20L).atLeast(1)).releaseResource(eq(taskExecutorConnection.getInstanceID()), any(Exception.class));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID());

		assertEquals(0, slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that idle but not releasable task managers will not be released even if timed out before it can be.
 */
@Test
public void testTaskManagerNotReleasedBeforeItCanBe() throws Exception {
	final CompletableFuture<InstanceID> releaseFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releaseFuture.complete(instanceID))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();

	final AtomicBoolean canBeReleased = new AtomicBoolean(false);
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setCanBeReleasedSupplier(canBeReleased::get)
		.createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final ManuallyTriggeredScheduledExecutor mainThreadExecutor = new ManuallyTriggeredScheduledExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setScheduledExecutor(mainThreadExecutor)
		.setTaskManagerTimeout(Time.milliseconds(0L))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		mainThreadExecutor.execute(() -> slotManager.registerTaskManager(taskManagerConnection, slotReport));

		// now it can not be released yet
		canBeReleased.set(false);
		mainThreadExecutor.execute(slotManager::checkTaskManagerTimeouts);
		mainThreadExecutor.triggerAll();
		assertFalse(releaseFuture.isDone());

		// now it can and should be released
		canBeReleased.set(true);
		mainThreadExecutor.execute(slotManager::checkTaskManagerTimeouts);
		mainThreadExecutor.triggerAll();
		assertThat(releaseFuture.get(), is(equalTo(taskManagerConnection.getInstanceID())));
	}
}
 

@Test
public void testHeartbeatTimeoutWithResourceManager() throws Exception {
	final String resourceManagerAddress = "rm";
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID rmResourceId = new ResourceID(resourceManagerAddress);

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway(
		resourceManagerId,
		rmResourceId,
		resourceManagerAddress,
		"localhost");

	final CompletableFuture<Tuple3<JobMasterId, ResourceID, JobID>> jobManagerRegistrationFuture = new CompletableFuture<>();
	final CompletableFuture<JobID> disconnectedJobManagerFuture = new CompletableFuture<>();
	final CountDownLatch registrationAttempts = new CountDownLatch(2);

	resourceManagerGateway.setRegisterJobManagerConsumer(tuple -> {
		jobManagerRegistrationFuture.complete(
			Tuple3.of(
				tuple.f0,
				tuple.f1,
				tuple.f3));
		registrationAttempts.countDown();
	});

	resourceManagerGateway.setDisconnectJobManagerConsumer(tuple -> disconnectedJobManagerFuture.complete(tuple.f0));

	rpcService.registerGateway(resourceManagerAddress, resourceManagerGateway);

	final JobManagerSharedServices jobManagerSharedServices = new TestingJobManagerSharedServicesBuilder().build();

	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		jobManagerSharedServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start operation to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// define a leader and see that a registration happens
		rmLeaderRetrievalService.notifyListener(resourceManagerAddress, resourceManagerId.toUUID());

		// register job manager success will trigger monitor heartbeat target between jm and rm
		final Tuple3<JobMasterId, ResourceID, JobID> registrationInformation = jobManagerRegistrationFuture.get(
			testingTimeout.toMilliseconds(),
			TimeUnit.MILLISECONDS);

		assertThat(registrationInformation.f0, Matchers.equalTo(jobMasterId));
		assertThat(registrationInformation.f1, Matchers.equalTo(jmResourceId));
		assertThat(registrationInformation.f2, Matchers.equalTo(jobGraph.getJobID()));

		final JobID disconnectedJobManager = disconnectedJobManagerFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// heartbeat timeout should trigger disconnect JobManager from ResourceManager
		assertThat(disconnectedJobManager, Matchers.equalTo(jobGraph.getJobID()));

		// the JobMaster should try to reconnect to the RM
		registrationAttempts.await();
	} finally {
		jobManagerSharedServices.shutdown();
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * 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()
		.setScheduledExecutor(scheduledExecutor)
		.build()) {

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

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

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

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

		slotManager.registerSlotRequest(slotRequest);

		assertThat(resourceProfileFuture.get(), is(equalTo(slotRequest.getResourceProfile())));

		// 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 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 = new ResourceProfile(1.0, 1);
	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");

	final Executor mainThreadExecutor = TestingUtils.defaultExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.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();
		}
	}
}
 

/**
 * 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 ResourceID resourceID = ResourceID.generate();
	final JobID jobId = new JobID();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);

	ResourceActions resourceManagerActions = mock(ResourceActions.class);

	// accept an incoming slot request
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

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

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

	try (SlotManager 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());
	}
}
 

/**
 * Tests that pending slot requests are tried to be fulfilled upon new slot registrations.
 */
@Test
public void testFulfillingPendingSlotRequest() 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 = new ResourceProfile(42.0, 1337);
	final SlotRequest slotRequest = new SlotRequest(
		jobId,
		allocationId,
		resourceProfile,
		targetAddress);

	final AtomicInteger numberAllocateResourceCalls = new AtomicInteger(0);
	ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceConsumer(ignored -> numberAllocateResourceCalls.incrementAndGet())
		.build();

	final CompletableFuture<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>> requestFuture = new CompletableFuture<>();
	// accept an incoming slot request
	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(tuple5 -> {
			requestFuture.complete(Tuple5.of(tuple5.f0, tuple5.f1, tuple5.f2, tuple5.f3, tuple5.f4));
			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);

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

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

		assertThat(numberAllocateResourceCalls.get(), is(1));

		slotManager.registerTaskManager(
			taskExecutorConnection,
			slotReport);

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

		TaskManagerSlot slot = slotManager.getSlot(slotId);

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

/**
 * 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()
		.setScheduledExecutor(scheduledExecutor)
		.build()) {

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

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

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

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

		slotManager.registerSlotRequest(slotRequest);

		assertThat(resourceProfileFuture.get(), is(equalTo(slotRequest.getResourceProfile())));

		// 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 un-registration of task managers will free and remove all registered slots.
 */
@Test
public void testTaskManagerUnregistration() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final JobID jobId = new JobID();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(tuple5 -> {
			assertThat(tuple5.f4, is(equalTo(resourceManagerId)));
			return new CompletableFuture<>();
		})
		.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 AllocationID allocationId1 = new AllocationID();
	final AllocationID allocationId2 = new AllocationID();
	final ResourceProfile resourceProfile = new ResourceProfile(42.0, 1337);
	final SlotStatus slotStatus1 = new SlotStatus(slotId1, resourceProfile, jobId, allocationId1);
	final SlotStatus slotStatus2 = new SlotStatus(slotId2, resourceProfile);
	final SlotReport slotReport = new SlotReport(Arrays.asList(slotStatus1, slotStatus2));

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

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

		assertTrue("The number registered slots does not equal the expected number.", 2 == slotManager.getNumberRegisteredSlots());

		TaskManagerSlot slot1 = slotManager.getSlot(slotId1);
		TaskManagerSlot slot2 = slotManager.getSlot(slotId2);

		assertTrue(slot1.getState() == TaskManagerSlot.State.ALLOCATED);
		assertTrue(slot2.getState() == TaskManagerSlot.State.FREE);

		assertTrue(slotManager.registerSlotRequest(slotRequest));

		assertFalse(slot2.getState() == TaskManagerSlot.State.FREE);
		assertTrue(slot2.getState() == TaskManagerSlot.State.PENDING);

		PendingSlotRequest pendingSlotRequest = slotManager.getSlotRequest(allocationId2);

		assertTrue("The pending slot request should have been assigned to slot 2", pendingSlotRequest.isAssigned());

		slotManager.unregisterTaskManager(taskManagerConnection.getInstanceID());

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
		assertFalse(pendingSlotRequest.isAssigned());
	}
}
 

/**
 * 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 = new ResourceProfile(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<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>> requestFuture = new CompletableFuture<>();
		// accept an incoming slot request
		final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setRequestSlotFunction(tuple5 -> {
				requestFuture.complete(Tuple5.of(tuple5.f0, tuple5.f1, tuple5.f2, tuple5.f3, tuple5.f4));
				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(Tuple5.of(slotId, jobId, allocationId, targetAddress, resourceManagerId))));

		TaskManagerSlot slot = slotManager.getSlot(slotId);

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