Java Code Examples for org.apache.flink.runtime.registration.RetryingRegistrationConfiguration#defaultConfiguration()

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

public TaskManagerServicesBuilder() {
	taskManagerLocation = new LocalTaskManagerLocation();
	memoryManager = new MemoryManager(
		MemoryManager.MIN_PAGE_SIZE,
		1,
		MemoryManager.MIN_PAGE_SIZE,
		MemoryType.HEAP,
		false);
	ioManager = mock(IOManager.class);
	shuffleEnvironment = mock(ShuffleEnvironment.class);
	kvStateService = new KvStateService(new KvStateRegistry(), null, null);
	broadcastVariableManager = new BroadcastVariableManager();
	taskEventDispatcher = new TaskEventDispatcher();
	taskSlotTable = mock(TaskSlotTable.class);
	jobManagerTable = new JobManagerTable();
	jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());
	taskStateManager = mock(TaskExecutorLocalStateStoresManager.class);
}
 

private TaskExecutorToResourceManagerConnection createTaskExecutorToResourceManagerConnection() {
	final TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
		TASK_MANAGER_ADDRESS,
		TASK_MANAGER_RESOURCE_ID,
		TASK_MANAGER_DATA_PORT,
		TASK_MANAGER_HARDWARE_DESCRIPTION,
		ResourceProfile.ZERO,
		ResourceProfile.ZERO
	);
	return new TaskExecutorToResourceManagerConnection(
		LOGGER,
		rpcService,
		RetryingRegistrationConfiguration.defaultConfiguration(),
		RESOURCE_MANAGER_ADDRESS,
		RESOURCE_MANAGER_ID,
		Executors.directExecutor(),
		new TestRegistrationConnectionListener<>(),
		taskExecutorRegistration);
}
 

private TaskExecutorTestingContext createTaskExecutorTestingContext(final TaskSlotTable<Task> taskSlotTable) throws IOException {
	final OneShotLatch offerSlotsLatch = new OneShotLatch();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
		.setOfferSlotsFunction((resourceID, slotOffers) -> {
			offerSlotsLatch.trigger();
			return CompletableFuture.completedFuture(slotOffers);
		}).build();
	rpc.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);

	final JobLeaderService jobLeaderService = new DefaultJobLeaderService(
		unresolvedTaskManagerLocation,
		RetryingRegistrationConfiguration.defaultConfiguration());

	TaskExecutorLocalStateStoresManager stateStoresManager = createTaskExecutorLocalStateStoresManager();
	final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(stateStoresManager)
		.build());

	jobManagerLeaderRetriever.notifyListener(jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());
	return new TaskExecutorTestingContext(jobMasterGateway, taskSlotTable, taskExecutor);
}
 

private TaskExecutorToResourceManagerConnection createTaskExecutorToResourceManagerConnection() {
	return new TaskExecutorToResourceManagerConnection(
		LOGGER,
		rpcService,
		TASK_MANAGER_ADDRESS,
		TASK_MANAGER_RESOURCE_ID,
		RetryingRegistrationConfiguration.defaultConfiguration(),
		TASK_MANAGER_DATA_PORT,
		TASK_MANAGER_HARDWARE_DESCRIPTION,
		RESOURCE_MANAGER_ADDRESS,
		RESOURCE_MANAGER_ID,
		Executors.directExecutor(),
		new TestRegistrationConnectionListener<>());
}
 

/**
 * Returns a task manager services configuration for the tests
 *
 * @param managedMemory         see {@link TaskManagerOptions#MANAGED_MEMORY_SIZE}
 * @param managedMemoryFraction see {@link TaskManagerOptions#MANAGED_MEMORY_FRACTION}
 * @param networkBufFraction	see {@link TaskManagerOptions#NETWORK_BUFFERS_MEMORY_FRACTION}
 * @param networkBufMin			see {@link TaskManagerOptions#NETWORK_BUFFERS_MEMORY_MIN}
 * @param networkBufMax			see {@link TaskManagerOptions#NETWORK_BUFFERS_MEMORY_MAX}
 * @param memType				on-heap or off-heap
 *
 * @return configuration object
 */
private static TaskManagerServicesConfiguration getTmConfig(
	final long managedMemory, final float managedMemoryFraction, float networkBufFraction,
	long networkBufMin, long networkBufMax,
	final MemoryType memType) {

	final NetworkEnvironmentConfiguration networkConfig = new NetworkEnvironmentConfiguration(
		networkBufFraction,
		networkBufMin,
		networkBufMax,
		checkedDownCast(MemorySize.parse(TaskManagerOptions.MEMORY_SEGMENT_SIZE.defaultValue()).getBytes()),
		null,
		TaskManagerOptions.NETWORK_REQUEST_BACKOFF_INITIAL.defaultValue(),
		TaskManagerOptions.NETWORK_REQUEST_BACKOFF_MAX.defaultValue(),
		TaskManagerOptions.NETWORK_BUFFERS_PER_CHANNEL.defaultValue(),
		TaskManagerOptions.NETWORK_EXTRA_BUFFERS_PER_GATE.defaultValue(),
		null);

	return new TaskManagerServicesConfiguration(
		InetAddress.getLoopbackAddress(),
		new String[] {},
		new String[] {},
		false,
		networkConfig,
		QueryableStateConfiguration.disabled(),
		1,
		managedMemory,
		memType,
		false,
		managedMemoryFraction,
		0,
		RetryingRegistrationConfiguration.defaultConfiguration(),
		Optional.empty());
}
 

private TaskExecutorToResourceManagerConnection createTaskExecutorToResourceManagerConnection() {
	return new TaskExecutorToResourceManagerConnection(
		LOGGER,
		rpcService,
		TASK_MANAGER_ADDRESS,
		TASK_MANAGER_RESOURCE_ID,
		RetryingRegistrationConfiguration.defaultConfiguration(),
		TASK_MANAGER_DATA_PORT,
		TASK_MANAGER_HARDWARE_DESCRIPTION,
		RESOURCE_MANAGER_ADDRESS,
		RESOURCE_MANAGER_ID,
		Executors.directExecutor(),
		new TestRegistrationConnectionListener<>());
}
 

private JobLeaderService createAndStartJobLeaderService(HighAvailabilityServices haServices, JobLeaderListener testingJobLeaderListener) {
	final JobLeaderService jobLeaderService = new DefaultJobLeaderService(
		new LocalUnresolvedTaskManagerLocation(),
		RetryingRegistrationConfiguration.defaultConfiguration());

	jobLeaderService.start(
		"foobar",
		rpcServiceResource.getTestingRpcService(),
		haServices,
		testingJobLeaderListener);
	return jobLeaderService;
}
 

public TaskManagerServicesBuilder() {
	unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
	ioManager = mock(IOManager.class);
	shuffleEnvironment = mock(ShuffleEnvironment.class);
	kvStateService = new KvStateService(new KvStateRegistry(), null, null);
	broadcastVariableManager = new BroadcastVariableManager();
	taskEventDispatcher = new TaskEventDispatcher();
	taskSlotTable = TestingTaskSlotTable.<Task>newBuilder().closeAsyncReturns(CompletableFuture.completedFuture(null)).build();
	jobTable = DefaultJobTable.create();
	jobLeaderService = new DefaultJobLeaderService(unresolvedTaskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());
	taskStateManager = mock(TaskExecutorLocalStateStoresManager.class);
	ioExecutor = TestingUtils.defaultExecutor();
	libraryCacheManager = TestingLibraryCacheManager.newBuilder().build();
}
 

@Test
public void testShouldShutDownTaskManagerServicesInPostStop() throws Exception {
	final TaskSlotTableImpl<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);

	final JobLeaderService jobLeaderService = new DefaultJobLeaderService(unresolvedTaskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final IOManager ioManager = new IOManagerAsync(tmp.newFolder().getAbsolutePath());

	final TaskExecutorLocalStateStoresManager localStateStoresManager = new TaskExecutorLocalStateStoresManager(
		false,
		ioManager.getSpillingDirectories(),
		Executors.directExecutor());

	nettyShuffleEnvironment.start();

	final KvStateService kvStateService = new KvStateService(new KvStateRegistry(), null, null);
	kvStateService.start();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation)
		.setIoManager(ioManager)
		.setShuffleEnvironment(nettyShuffleEnvironment)
		.setKvStateService(kvStateService)
		.setTaskSlotTable(taskSlotTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}

	assertThat(taskSlotTable.isClosed(), is(true));
	assertThat(nettyShuffleEnvironment.isClosed(), is(true));
	assertThat(kvStateService.isShutdown(), is(true));
}
 

@Test
public void testShouldShutDownTaskManagerServicesInPostStop() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);

	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final IOManager ioManager = new IOManagerAsync(tmp.newFolder().getAbsolutePath());

	final TaskExecutorLocalStateStoresManager localStateStoresManager = new TaskExecutorLocalStateStoresManager(
		false,
		ioManager.getSpillingDirectories(),
		Executors.directExecutor());

	final MemoryManager memoryManager = new MemoryManager(
		4096,
		1,
		4096,
		MemoryType.HEAP,
		false);

	final NetworkEnvironment networkEnvironment = new NetworkEnvironment(
		1,
		1,
		0,
		0,
		2,
		8,
		true);
	networkEnvironment.start();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setMemoryManager(memoryManager)
		.setIoManager(ioManager)
		.setNetworkEnvironment(networkEnvironment)
		.setTaskSlotTable(taskSlotTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final long heartbeatInterval = 1000L;
	final long heartbeatTimeout = 1000L;
	final HeartbeatServices heartbeatServices = new HeartbeatServices(heartbeatInterval, heartbeatTimeout);

	final TaskExecutor taskManager = new TaskExecutor(
		rpc,
		taskManagerConfiguration,
		haServices,
		taskManagerServices,
		heartbeatServices,
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	try {
		taskManager.start();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}

	assertThat(memoryManager.isShutdown(), is(true));
	assertThat(networkEnvironment.isShutdown(), is(true));
	assertThat(ioManager.isProperlyShutDown(), is(true));
}
 

/**
 * Tests that a TaskManager detects a job leader for which it has reserved slots. Upon detecting
 * the job leader, it will offer all reserved slots to the JobManager.
 */
@Test
public void testJobLeaderDetection() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final JobManagerTable jobManagerTable = new JobManagerTable();
	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	CompletableFuture<Void> initialSlotReportFuture = new CompletableFuture<>();
	resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<Collection<SlotOffer>> offeredSlotsFuture = new CompletableFuture<>();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
		.setOfferSlotsFunction((resourceID, slotOffers) -> {

			offeredSlotsFuture.complete(new ArrayList<>(slotOffers));
			return CompletableFuture.completedFuture(slotOffers);
		})
		.build();

	rpc.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
	rpc.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);

	final AllocationID allocationId = new AllocationID();
	final SlotID slotId = new SlotID(taskManagerLocation.getResourceID(), 0);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setJobManagerTable(jobManagerTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = new TaskExecutor(
		rpc,
		taskManagerConfiguration,
		haServices,
		taskManagerServices,
		HEARTBEAT_SERVICES,
		UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
		null,
		dummyBlobCacheService,
		testingFatalErrorHandler);

	try {
		taskManager.start();

		final TaskExecutorGateway tmGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// tell the task manager about the rm leader
		resourceManagerLeaderRetriever.notifyListener(resourceManagerGateway.getAddress(), resourceManagerGateway.getFencingToken().toUUID());

		// wait for the initial slot report
		initialSlotReportFuture.get();

		// request slots from the task manager under the given allocation id
		CompletableFuture<Acknowledge> slotRequestAck = tmGateway.requestSlot(
			slotId,
			jobId,
			allocationId,
			jobMasterGateway.getAddress(),
			resourceManagerGateway.getFencingToken(),
			timeout);

		slotRequestAck.get();

		// now inform the task manager about the new job leader
		jobManagerLeaderRetriever.notifyListener(jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());

		final Collection<SlotOffer> offeredSlots = offeredSlotsFuture.get();
		final Collection<AllocationID> allocationIds = offeredSlots.stream().map(SlotOffer::getAllocationId).collect(Collectors.toList());
		assertThat(allocationIds, containsInAnyOrder(allocationId));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

@Test
public void testShouldShutDownTaskManagerServicesInPostStop() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);

	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final IOManager ioManager = new IOManagerAsync(tmp.newFolder().getAbsolutePath());

	final TaskExecutorLocalStateStoresManager localStateStoresManager = new TaskExecutorLocalStateStoresManager(
		false,
		ioManager.getSpillingDirectories(),
		Executors.directExecutor());

	final MemoryManager memoryManager = new MemoryManager(
		4096,
		1,
		4096,
		MemoryType.HEAP,
		false);

	nettyShuffleEnvironment.start();

	final KvStateService kvStateService = new KvStateService(new KvStateRegistry(), null, null);
	kvStateService.start();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setMemoryManager(memoryManager)
		.setIoManager(ioManager)
		.setShuffleEnvironment(nettyShuffleEnvironment)
		.setKvStateService(kvStateService)
		.setTaskSlotTable(taskSlotTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}

	assertThat(memoryManager.isShutdown(), is(true));
	assertThat(nettyShuffleEnvironment.isClosed(), is(true));
	assertThat(kvStateService.isShutdown(), is(true));
}
 

/**
 * Tests that a TaskManager detects a job leader for which it has reserved slots. Upon detecting
 * the job leader, it will offer all reserved slots to the JobManager.
 */
@Test
public void testJobLeaderDetection() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final JobManagerTable jobManagerTable = new JobManagerTable();
	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	CompletableFuture<Void> initialSlotReportFuture = new CompletableFuture<>();
	resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<Collection<SlotOffer>> offeredSlotsFuture = new CompletableFuture<>();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
		.setOfferSlotsFunction((resourceID, slotOffers) -> {

			offeredSlotsFuture.complete(new ArrayList<>(slotOffers));
			return CompletableFuture.completedFuture(slotOffers);
		})
		.build();

	rpc.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
	rpc.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);

	final AllocationID allocationId = new AllocationID();
	final SlotID slotId = new SlotID(taskManagerLocation.getResourceID(), 0);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setJobManagerTable(jobManagerTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		final TaskExecutorGateway tmGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// tell the task manager about the rm leader
		resourceManagerLeaderRetriever.notifyListener(resourceManagerGateway.getAddress(), resourceManagerGateway.getFencingToken().toUUID());

		// wait for the initial slot report
		initialSlotReportFuture.get();

		// request slots from the task manager under the given allocation id
		CompletableFuture<Acknowledge> slotRequestAck = tmGateway.requestSlot(
			slotId,
			jobId,
			allocationId,
			jobMasterGateway.getAddress(),
			resourceManagerGateway.getFencingToken(),
			timeout);

		slotRequestAck.get();

		// now inform the task manager about the new job leader
		jobManagerLeaderRetriever.notifyListener(jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());

		final Collection<SlotOffer> offeredSlots = offeredSlotsFuture.get();
		final Collection<AllocationID> allocationIds = offeredSlots.stream().map(SlotOffer::getAllocationId).collect(Collectors.toList());
		assertThat(allocationIds, containsInAnyOrder(allocationId));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that accepted slots go into state assigned and the others are returned to the resource
 * manager.
 */
@Test
public void testSlotAcceptance() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Arrays.asList(mock(ResourceProfile.class), mock(ResourceProfile.class)), timerService);
	final JobManagerTable jobManagerTable = new JobManagerTable();
	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final String resourceManagerAddress = "rm";
	final UUID resourceManagerLeaderId = UUID.randomUUID();

	final String jobManagerAddress = "jm";
	final UUID jobManagerLeaderId = UUID.randomUUID();

	resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, resourceManagerLeaderId);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobManagerLeaderId);

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	final ResourceID resourceManagerResourceId = resourceManagerGateway.getOwnResourceId();
	final InstanceID registrationId = new InstanceID();

	final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
	resourceManagerGateway.setRegisterTaskExecutorFunction(
		stringResourceIDIntegerHardwareDescriptionTuple4 -> {
			registrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
			return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, new ClusterInformation("localhost", 1234)));
		});

	final CompletableFuture<Tuple3<InstanceID, SlotID, AllocationID>> availableSlotFuture = new CompletableFuture<>();
	resourceManagerGateway.setNotifySlotAvailableConsumer(availableSlotFuture::complete);

	final ResourceID jmResourceId = new ResourceID(jobManagerAddress);

	final AllocationID allocationId1 = new AllocationID();
	final AllocationID allocationId2 = new AllocationID();

	final SlotOffer offer1 = new SlotOffer(allocationId1, 0, ResourceProfile.UNKNOWN);

	final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);

	when(jobMasterGateway.registerTaskManager(
			any(String.class),
			eq(taskManagerLocation),
			any(Time.class)
	)).thenReturn(CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jmResourceId)));
	when(jobMasterGateway.getHostname()).thenReturn(jobManagerAddress);

	when(jobMasterGateway.offerSlots(
			any(ResourceID.class), any(Collection.class), any(Time.class)))
		.thenReturn(CompletableFuture.completedFuture((Collection<SlotOffer>) Collections.singleton(offer1)));

	rpc.registerGateway(resourceManagerAddress, resourceManagerGateway);
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskManagerLocation(taskManagerLocation)
		.setTaskSlotTable(taskSlotTable)
		.setJobManagerTable(jobManagerTable)
		.setJobLeaderService(jobLeaderService)
		.setTaskStateManager(localStateStoresManager)
		.build();

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		assertThat(registrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));

		taskSlotTable.allocateSlot(0, jobId, allocationId1, Time.milliseconds(10000L));
		taskSlotTable.allocateSlot(1, jobId, allocationId2, Time.milliseconds(10000L));

		// we have to add the job after the TaskExecutor, because otherwise the service has not
		// been properly started.
		jobLeaderService.addJob(jobId, jobManagerAddress);

		final Tuple3<InstanceID, SlotID, AllocationID> instanceIDSlotIDAllocationIDTuple3 = availableSlotFuture.get();

		final Tuple3<InstanceID, SlotID, AllocationID> expectedResult = Tuple3.of(registrationId, new SlotID(taskManagerLocation.getResourceID(), 1), allocationId2);

		assertThat(instanceIDSlotIDAllocationIDTuple3, equalTo(expectedResult));

		assertTrue(taskSlotTable.tryMarkSlotActive(jobId, allocationId1));
		assertFalse(taskSlotTable.tryMarkSlotActive(jobId, allocationId2));
		assertTrue(taskSlotTable.isSlotFree(1));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}
 

/**
 * Tests that a TaskManager detects a job leader for which it has reserved slots. Upon detecting
 * the job leader, it will offer all reserved slots to the JobManager.
 */
@Test
public void testJobLeaderDetection() throws Exception {
	final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
	final JobLeaderService jobLeaderService = new DefaultJobLeaderService(unresolvedTaskManagerLocation, RetryingRegistrationConfiguration.defaultConfiguration());

	final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
	CompletableFuture<Void> initialSlotReportFuture = new CompletableFuture<>();
	resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReportFuture.complete(null);
		return CompletableFuture.completedFuture(Acknowledge.get());
	});

	final CompletableFuture<Collection<SlotOffer>> offeredSlotsFuture = new CompletableFuture<>();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
		.setOfferSlotsFunction((resourceID, slotOffers) -> {

			offeredSlotsFuture.complete(new ArrayList<>(slotOffers));
			return CompletableFuture.completedFuture(slotOffers);
		})
		.build();

	rpc.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
	rpc.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);

	final AllocationID allocationId = new AllocationID();
	final SlotID slotId = new SlotID(unresolvedTaskManagerLocation.getResourceID(), 0);

	final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();

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

	TaskExecutor taskManager = createTaskExecutor(taskManagerServices);

	try {
		taskManager.start();

		final TaskExecutorGateway tmGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);

		// tell the task manager about the rm leader
		resourceManagerLeaderRetriever.notifyListener(resourceManagerGateway.getAddress(), resourceManagerGateway.getFencingToken().toUUID());

		// wait for the initial slot report
		initialSlotReportFuture.get();

		// request slots from the task manager under the given allocation id
		CompletableFuture<Acknowledge> slotRequestAck = tmGateway.requestSlot(
			slotId,
			jobId,
			allocationId,
			ResourceProfile.ZERO,
			jobMasterGateway.getAddress(),
			resourceManagerGateway.getFencingToken(),
			timeout);

		slotRequestAck.get();

		// now inform the task manager about the new job leader
		jobManagerLeaderRetriever.notifyListener(jobMasterGateway.getAddress(), jobMasterGateway.getFencingToken().toUUID());

		final Collection<SlotOffer> offeredSlots = offeredSlotsFuture.get();
		final Collection<AllocationID> allocationIds = offeredSlots.stream().map(SlotOffer::getAllocationId).collect(Collectors.toList());
		assertThat(allocationIds, containsInAnyOrder(allocationId));
	} finally {
		RpcUtils.terminateRpcEndpoint(taskManager, timeout);
	}
}