org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway Java Examples

private void startDispatcher(JobManagerRunnerFactory jobManagerRunnerFactory) throws Exception {
	TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	final HeartbeatServices heartbeatServices = new HeartbeatServices(1000L, 1000L);
	final MemoryArchivedExecutionGraphStore archivedExecutionGraphStore = new MemoryArchivedExecutionGraphStore();
	dispatcher = new TestingDispatcher(
		rpcService,
		DispatcherId.generate(),
		new DefaultDispatcherBootstrap(Collections.emptyList()),
		new DispatcherServices(
			configuration,
			highAvailabilityServices,
			() -> CompletableFuture.completedFuture(resourceManagerGateway),
			blobServer,
			heartbeatServices,
			archivedExecutionGraphStore,
			testingFatalErrorHandlerResource.getFatalErrorHandler(),
			VoidHistoryServerArchivist.INSTANCE,
			null,
			UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup(),
			jobGraphWriter,
			jobManagerRunnerFactory));

	dispatcher.start();

	dispatcherGateway = dispatcher.getSelfGateway(DispatcherGateway.class);
}
 

/**
 * Tests that a batch slot request does react to {@link SlotPool#failAllocation(AllocationID, Exception)}
 * signals whose exception is {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestFailsIfAllocationFailsUnfulfillably() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	testingResourceManagerGateway.setRequestSlotConsumer(slotRequest -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		SlotPoolUtils.failAllocation(slotPool, directMainThreadExecutor, allocationIdFuture.get(),
			new UnfulfillableSlotRequestException(new AllocationID(), ResourceProfile.UNKNOWN));

		assertThat(slotFuture.isCompletedExceptionally(), is(true));
	}
}
 

@Nonnull
private HATestingDispatcher createDispatcher(
	HighAvailabilityServices highAvailabilityServices,
	@Nonnull Queue<DispatcherId> fencingTokens,
	JobManagerRunnerFactory jobManagerRunnerFactory) throws Exception {
	final Configuration configuration = new Configuration();

	TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	return new HATestingDispatcher(
		rpcService,
		UUID.randomUUID().toString(),
		configuration,
		highAvailabilityServices,
		() -> CompletableFuture.completedFuture(resourceManagerGateway),
		new BlobServer(configuration, new VoidBlobStore()),
		new HeartbeatServices(1000L, 1000L),
		UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup(),
		null,
		new MemoryArchivedExecutionGraphStore(),
		jobManagerRunnerFactory,
		testingFatalErrorHandler,
		fencingTokens);
}
 

/**
 * Tests that a batch slot request won't fail if its resource manager request fails with exceptions other than
 * {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestDoesNotFailIfRMRequestFails() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotFuture(FutureUtils.completedExceptionally(new FlinkException("Failed request")));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	final Time batchSlotTimeout = Time.milliseconds(1000L);
	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setBatchSlotTimeout(batchSlotTimeout)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		assertThat(slotFuture.isDone(), is(false));
	}
}
 

/**
 * Tests that the {@link TaskExecutor} sends the initial slot report after it
 * registered at the ResourceManager.
 */
@Test
public void testInitialSlotReport() throws Exception {
	final TaskExecutor taskExecutor = createTaskExecutor(1);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f0);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		assertThat(initialSlotReportFuture.get(), equalTo(taskExecutor.getResourceID()));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that a batch slot request fails if its resource manager request fails with {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestFailsIfRMRequestFailsUnfulfillably() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotFuture(FutureUtils.completedExceptionally(
		new UnfulfillableSlotRequestException(new AllocationID(), ResourceProfile.UNKNOWN)));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		assertThat(slotFuture.isCompletedExceptionally(), is(true));
	}
}
 

/**
 * Tests that a batch slot request does not react to {@link SlotPool#failAllocation(AllocationID, Exception)}
 * signals whose exception is not {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestDoesNotFailIfAllocationFails() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	testingResourceManagerGateway.setRequestSlotConsumer(slotRequest -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	final Time batchSlotTimeout = Time.milliseconds(1000L);
	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setBatchSlotTimeout(batchSlotTimeout)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		SlotPoolUtils.failAllocation(slotPool, directMainThreadExecutor, allocationIdFuture.get(), new FlinkException("Failed request"));

		assertThat(slotFuture.isDone(), is(false));
	}
}
 

@Nonnull
private TestingDispatcher createDispatcher(HeartbeatServices heartbeatServices, TestingHighAvailabilityServices haServices, JobManagerRunnerFactory jobManagerRunnerFactory) throws Exception {
	TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	return new TestingDispatcher(
		rpcService,
		Dispatcher.DISPATCHER_NAME + '_' + name.getMethodName(),
		configuration,
		haServices,
		() -> CompletableFuture.completedFuture(resourceManagerGateway),
		blobServer,
		heartbeatServices,
		UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup(),
		null,
		new MemoryArchivedExecutionGraphStore(),
		jobManagerRunnerFactory,
		fatalErrorHandler);
}
 

/**
 * Tests that a batch slot request does not react to {@link SlotPool#failAllocation(AllocationID, Exception)}
 * signals whose exception is not {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestDoesNotFailIfAllocationFails() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();
	testingResourceManagerGateway.setRequestSlotConsumer(slotRequest -> allocationIdFuture.complete(slotRequest.getAllocationId()));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	final Time batchSlotTimeout = Time.milliseconds(1000L);
	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setBatchSlotTimeout(batchSlotTimeout)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		SlotPoolUtils.failAllocation(slotPool, directMainThreadExecutor, allocationIdFuture.get(), new FlinkException("Failed request"));

		assertThat(slotFuture.isDone(), is(false));
	}
}
 

/**
 * Tests that a batch slot request won't fail if its resource manager request fails with exceptions other than
 * {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestDoesNotFailIfRMRequestFails() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotFuture(FutureUtils.completedExceptionally(new FlinkException("Failed request")));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	final Time batchSlotTimeout = Time.milliseconds(1000L);
	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setBatchSlotTimeout(batchSlotTimeout)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		assertThat(slotFuture.isDone(), is(false));
	}
}
 

/**
 * Tests that a batch slot request fails if its resource manager request fails with {@link UnfulfillableSlotRequestException}.
 */
@Test
public void testPendingBatchSlotRequestFailsIfRMRequestFailsUnfulfillably() throws Exception {
	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	testingResourceManagerGateway.setRequestSlotFuture(FutureUtils.completedExceptionally(
		new UnfulfillableSlotRequestException(new AllocationID(), ResourceProfile.UNKNOWN)));

	final ComponentMainThreadExecutor directMainThreadExecutor = ComponentMainThreadExecutorServiceAdapter.forMainThread();

	try (final SlotPoolImpl slotPool = new SlotPoolBuilder(directMainThreadExecutor)
		.setResourceManagerGateway(testingResourceManagerGateway)
		.build()) {

		final CompletableFuture<PhysicalSlot> slotFuture = SlotPoolUtils.requestNewAllocatedBatchSlot(slotPool, directMainThreadExecutor, resourceProfile);

		assertThat(slotFuture.isCompletedExceptionally(), is(true));
	}
}
 

private static Scheduler createSchedulerWithSlots(
		int numSlots,
		SlotPool slotPool,
		TaskManagerLocation taskManagerLocation) throws Exception {
	final TaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();
	final String jobManagerAddress = "foobar";
	final ResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	slotPool.start(JobMasterId.generate(), jobManagerAddress, mainThreadExecutor);
	slotPool.connectToResourceManager(resourceManagerGateway);
	Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), slotPool);
	scheduler.start(mainThreadExecutor);

	CompletableFuture.runAsync(() -> slotPool.registerTaskManager(taskManagerLocation.getResourceID()), mainThreadExecutor).join();

	final List<SlotOffer> slotOffers = new ArrayList<>(NUM_TASKS);
	for (int i = 0; i < numSlots; i++) {
		final AllocationID allocationId = new AllocationID();
		final SlotOffer slotOffer = new SlotOffer(allocationId, 0, ResourceProfile.ANY);
		slotOffers.add(slotOffer);
	}

	CompletableFuture.runAsync(() -> slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers), mainThreadExecutor).join();

	return scheduler;
}
 

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

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

@Test
public void testImmediatelyRegistersIfLeaderIsKnown() throws Exception {
	final String resourceManagerAddress = "/resource/manager/address/one";

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CountDownLatch taskManagerRegisteredLatch = new CountDownLatch(1);
	testingResourceManagerGateway.setRegisterTaskExecutorFunction(FunctionUtils.uncheckedFunction(
		ignored -> {
			taskManagerRegisteredLatch.countDown();
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
				new InstanceID(), new ResourceID(resourceManagerAddress), new ClusterInformation("localhost", 1234)));
		}
	));

	rpc.registerGateway(resourceManagerAddress, testingResourceManagerGateway);

	final TaskSlotTable taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, UUID.randomUUID());

		assertTrue(taskManagerRegisteredLatch.await(timeout.toMilliseconds(), TimeUnit.MILLISECONDS));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that the a JM connects to the leading RM after regaining leadership.
 */
@Test
public void testResourceManagerConnectionAfterRegainingLeadership() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

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

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

		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();

		final BlockingQueue<JobMasterId> registrationQueue = new ArrayBlockingQueue<>(1);
		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		final JobMasterId firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		jobMaster.suspend(new FlinkException("Test exception.")).get();

		final JobMasterId jobMasterId2 = JobMasterId.generate();

		jobMaster.start(jobMasterId2).get();

		final JobMasterId secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(jobMasterId2));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@Before
public void setUp() {
	rpcService = new TestingRpcService();

	testingResourceManagerGateway = new TestingResourceManagerGateway();
	rpcService.registerGateway(RESOURCE_MANAGER_ADDRESS, testingResourceManagerGateway);

	registrationSuccessFuture = new CompletableFuture<>();
}
 

@Test
public void testImmediatelyRegistersIfLeaderIsKnown() throws Exception {
	final String resourceManagerAddress = "/resource/manager/address/one";

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
	final CountDownLatch taskManagerRegisteredLatch = new CountDownLatch(1);
	testingResourceManagerGateway.setRegisterTaskExecutorFunction(FunctionUtils.uncheckedFunction(
		ignored -> {
			taskManagerRegisteredLatch.countDown();
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
				new InstanceID(), new ResourceID(resourceManagerAddress), new ClusterInformation("localhost", 1234)));
		}
	));

	rpc.registerGateway(resourceManagerAddress, testingResourceManagerGateway);

	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, UUID.randomUUID());

		assertTrue(taskManagerRegisteredLatch.await(timeout.toMilliseconds(), TimeUnit.MILLISECONDS));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

private CompletableFuture<RegistrationResponse> createRegistrationResponse(TestingResourceManagerGateway rmGateway1) {
	return CompletableFuture.completedFuture(
		new TaskExecutorRegistrationSuccess(
			new InstanceID(),
			rmGateway1.getOwnResourceId(),
			new ClusterInformation("localhost", 1234)));
}
 

/**
 * Tests that the {@link TaskExecutor} sends the initial slot report after it
 * registered at the ResourceManager.
 */
@Test
public void testInitialSlotReport() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f0);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		assertThat(initialSlotReportFuture.get(), equalTo(taskManagerLocation.getResourceID()));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the a JM connects to the leading RM after regaining leadership.
 */
@Test
public void testResourceManagerConnectionAfterRegainingLeadership() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

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

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

		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();

		final BlockingQueue<JobMasterId> registrationQueue = new ArrayBlockingQueue<>(1);
		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		final JobMasterId firstRegistrationAttempt = registrationQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		jobMaster.suspend(new FlinkException("Test exception.")).get();

		final JobMasterId jobMasterId2 = JobMasterId.generate();

		jobMaster.start(jobMasterId2).get();

		final JobMasterId secondRegistrationAttempt = registrationQueue.take();

		assertThat(secondRegistrationAttempt, equalTo(jobMasterId2));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Tests that we continue reconnecting to the latest known RM after a disconnection
 * message.
 */
@Test
public void testReconnectionAfterDisconnect() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

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

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
		final BlockingQueue<JobMasterId> registrationsQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationsQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		final ResourceManagerId resourceManagerId = testingResourceManagerGateway.getFencingToken();
		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		// wait for first registration attempt
		final JobMasterId firstRegistrationAttempt = registrationsQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		assertThat(registrationsQueue.isEmpty(), is(true));
		jobMasterGateway.disconnectResourceManager(resourceManagerId, new FlinkException("Test exception"));

		// wait for the second registration attempt after the disconnect call
		assertThat(registrationsQueue.take(), equalTo(jobMasterId));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Tests that we can close an unestablished ResourceManager connection.
 */
@Test
public void testCloseUnestablishedResourceManagerConnection() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	try {
		jobMaster.start(JobMasterId.generate()).get();

		final TestingResourceManagerGateway firstResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
		final TestingResourceManagerGateway secondResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();

		final OneShotLatch firstJobManagerRegistration = new OneShotLatch();
		final OneShotLatch secondJobManagerRegistration = new OneShotLatch();

		firstResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> firstJobManagerRegistration.trigger());

		secondResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> secondJobManagerRegistration.trigger());

		notifyResourceManagerLeaderListeners(firstResourceManagerGateway);

		// wait until we have seen the first registration attempt
		firstJobManagerRegistration.await();

		// this should stop the connection attempts towards the first RM
		notifyResourceManagerLeaderListeners(secondResourceManagerGateway);

		// check that we start registering at the second RM
		secondJobManagerRegistration.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

@Before
public void setUp() throws HandlerRequestException {
	resourceManagerGateway = new TestingResourceManagerGateway();
	taskManagerLogListHandler = new TaskManagerLogListHandler(
		() -> CompletableFuture.completedFuture(null),
		TestingUtils.TIMEOUT(),
		Collections.emptyMap(),
		TaskManagerLogsHeaders.getInstance(),
		() -> CompletableFuture.completedFuture(resourceManagerGateway));
	handlerRequest = createRequest(EXPECTED_TASK_MANAGER_ID);
}
 

/**
 * This case make sure when allocateSlot in ProviderAndOwner timeout,
 * it will automatically call cancelSlotAllocation as will inject future.whenComplete in ProviderAndOwner.
 */
@Test
public void testProviderAndOwnerSlotAllocationTimeout() throws Exception {
	final JobID jid = new JobID();

	try (TestingSlotPool pool = createTestingSlotPool(jid)) {

		final CompletableFuture<SlotRequestId> releaseSlotFuture = new CompletableFuture<>();

		pool.setReleaseSlotConsumer(releaseSlotFuture::complete);

		pool.start(JobMasterId.generate(), "foobar", testMainThreadExecutor.getMainThreadExecutor());
		ResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
		pool.connectToResourceManager(resourceManagerGateway);

		Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), pool);
		scheduler.start(testMainThreadExecutor.getMainThreadExecutor());

		// test the pending request is clear when timed out
		CompletableFuture<LogicalSlot> future = testMainThreadExecutor.execute(() -> scheduler.allocateSlot(
			new DummyScheduledUnit(),
			SlotProfile.noRequirements(),
			fastTimeout));
		try {
			future.get();
			fail("We expected a TimeoutException.");
		} catch (ExecutionException e) {
			assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
		}

		// wait for the cancel call on the SlotPoolImpl
		releaseSlotFuture.get();

		assertEquals(0L, pool.getNumberOfPendingRequests());
	}
}
 

@Override
protected void before() throws Throwable {
	if (slotPool != null) {
		terminateSlotPool();
	}

	testingResourceManagerGateway = new TestingResourceManagerGateway();

	slotPool = new TestingSlotPoolImpl(new JobID());
	scheduler = new SchedulerImpl(schedulingStrategy, slotPool);
	slotPool.start(JobMasterId.generate(), "foobar", mainThreadExecutor);
	scheduler.start(mainThreadExecutor);
	slotPool.connectToResourceManager(testingResourceManagerGateway);
}
 

@Override
protected void before() throws Throwable {
	if (slotPool != null) {
		terminateSlotPool();
	}

	testingResourceManagerGateway = new TestingResourceManagerGateway();

	slotPool = new SlotPoolImpl(new JobID());
	scheduler = new SchedulerImpl(schedulingStrategy, slotPool);
	slotPool.start(JobMasterId.generate(), "foobar", mainThreadExecutor);
	scheduler.start(mainThreadExecutor);
	slotPool.connectToResourceManager(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());
}
 

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

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

/**
 * This case make sure when allocateSlot in ProviderAndOwner timeout,
 * it will automatically call cancelSlotAllocation as will inject future.whenComplete in ProviderAndOwner.
 */
@Test
public void testProviderAndOwnerSlotAllocationTimeout() throws Exception {
	final JobID jid = new JobID();

	try (TestingSlotPool pool = createTestingSlotPool(jid)) {

		final CompletableFuture<SlotRequestId> releaseSlotFuture = new CompletableFuture<>();

		pool.setReleaseSlotConsumer(releaseSlotFuture::complete);

		pool.start(JobMasterId.generate(), "foobar", testMainThreadExecutor.getMainThreadExecutor());
		ResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
		pool.connectToResourceManager(resourceManagerGateway);

		Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.INSTANCE, pool);
		scheduler.start(testMainThreadExecutor.getMainThreadExecutor());

		// test the pending request is clear when timed out
		CompletableFuture<LogicalSlot> future = testMainThreadExecutor.execute(() -> scheduler.allocateSlot(
			new DummyScheduledUnit(),
			true,
			SlotProfile.noRequirements(),
			fastTimeout));
		try {
			future.get();
			fail("We expected a TimeoutException.");
		} catch (ExecutionException e) {
			assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
		}

		// wait for the cancel call on the SlotPoolImpl
		releaseSlotFuture.get();

		assertEquals(0L, pool.getNumberOfPendingRequests());
	}
}
 

/**
 * Tests that the {@link TaskExecutor} sends the initial slot report after it
 * registered at the ResourceManager.
 */
@Test
public void testInitialSlotReport() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
		final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();

		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f0);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

		rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
		resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

		assertThat(initialSlotReportFuture.get(), equalTo(taskManagerLocation.getResourceID()));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}