org.apache.flink.util.function.FunctionUtils Java Examples

private CompletableFuture<Void> runJob(JobGraph jobGraph) {
	Preconditions.checkState(!jobManagerRunnerFutures.containsKey(jobGraph.getJobID()));

	final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = createJobManagerRunner(jobGraph);

	jobManagerRunnerFutures.put(jobGraph.getJobID(), jobManagerRunnerFuture);

	return jobManagerRunnerFuture
		.thenApply(FunctionUtils.nullFn())
		.whenCompleteAsync(
			(ignored, throwable) -> {
				if (throwable != null) {
					jobManagerRunnerFutures.remove(jobGraph.getJobID());
				}
			},
			getMainThreadExecutor());
}
 

private CompletableFuture<JobManagerRunner> createJobManagerRunner(JobGraph jobGraph) {
	final RpcService rpcService = getRpcService();

	final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = CompletableFuture.supplyAsync(
		CheckedSupplier.unchecked(() ->
			jobManagerRunnerFactory.createJobManagerRunner(
				jobGraph,
				configuration,
				rpcService,
				highAvailabilityServices,
				heartbeatServices,
				jobManagerSharedServices,
				new DefaultJobManagerJobMetricGroupFactory(jobManagerMetricGroup),
				fatalErrorHandler)),
		rpcService.getExecutor());

	return jobManagerRunnerFuture.thenApply(FunctionUtils.uncheckedFunction(this::startJobManagerRunner));
}
 

private CompletableFuture<Void> runJob(JobGraph jobGraph) {
	Preconditions.checkState(!jobManagerRunnerFutures.containsKey(jobGraph.getJobID()));

	final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = createJobManagerRunner(jobGraph);

	jobManagerRunnerFutures.put(jobGraph.getJobID(), jobManagerRunnerFuture);

	return jobManagerRunnerFuture
		.thenApply(FunctionUtils.uncheckedFunction(this::startJobManagerRunner))
		.thenApply(FunctionUtils.nullFn())
		.whenCompleteAsync(
			(ignored, throwable) -> {
				if (throwable != null) {
					jobManagerRunnerFutures.remove(jobGraph.getJobID());
				}
			},
			getMainThreadExecutor());
}
 

private CompletableFuture<Void> runJob(JobGraph jobGraph) {
	Preconditions.checkState(!jobManagerRunnerFutures.containsKey(jobGraph.getJobID()));

	final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = createJobManagerRunner(jobGraph);

	jobManagerRunnerFutures.put(jobGraph.getJobID(), jobManagerRunnerFuture);

	return jobManagerRunnerFuture
		.thenApply(FunctionUtils.nullFn())
		.whenCompleteAsync(
			(ignored, throwable) -> {
				if (throwable != null) {
					jobManagerRunnerFutures.remove(jobGraph.getJobID());
				}
			},
			getMainThreadExecutor());
}
 

private CompletableFuture<JobManagerRunner> createJobManagerRunner(JobGraph jobGraph) {
	final RpcService rpcService = getRpcService();

	final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = CompletableFuture.supplyAsync(
		CheckedSupplier.unchecked(() ->
			jobManagerRunnerFactory.createJobManagerRunner(
				jobGraph,
				configuration,
				rpcService,
				highAvailabilityServices,
				heartbeatServices,
				jobManagerSharedServices,
				new DefaultJobManagerJobMetricGroupFactory(jobManagerMetricGroup),
				fatalErrorHandler)),
		rpcService.getExecutor());

	return jobManagerRunnerFuture.thenApply(FunctionUtils.uncheckedFunction(this::startJobManagerRunner));
}
 

public static List<File> checkJarFileForApplicationMode(Configuration configuration) {
	return configuration.get(PipelineOptions.JARS).stream().map(
		FunctionUtils.uncheckedFunction(
			uri -> {
				final URI jarURI = PackagedProgramUtils.resolveURI(uri);
				if (jarURI.getScheme().equals("local") && jarURI.isAbsolute()) {
					return new File(jarURI.getPath());
				}
				throw new IllegalArgumentException("Only \"local\" is supported as schema for application mode." +
						" This assumes that the jar is located in the image, not the Flink client." +
						" An example of such path is: local:///opt/flink/examples/streaming/WindowJoin.jar");
			})
	).collect(Collectors.toList());
}
 

@BeforeClass
public static void init() throws IOException {
	final String textFileName = "test.txt";
	final String userDirHasEntryClassName = "_test_user_dir_has_entry_class";
	final String userDirHasNotEntryClassName = "_test_user_dir_has_not_entry_class";

	userDirHasEntryClass = JOB_DIRS.newFolder(userDirHasEntryClassName);
	final Path userJarPath = userDirHasEntryClass.toPath().resolve(TestJobInfo.JOB_JAR_PATH.toFile().getName());
	final Path userLibJarPath =
		userDirHasEntryClass.toPath().resolve(TestJobInfo.JOB_LIB_JAR_PATH.toFile().getName());
	userDirHasNotEntryClass = JOB_DIRS.newFolder(userDirHasNotEntryClassName);

	//create files
	Files.copy(TestJobInfo.JOB_JAR_PATH, userJarPath);
	Files.copy(TestJobInfo.JOB_LIB_JAR_PATH, userLibJarPath);
	Files.createFile(userDirHasEntryClass.toPath().resolve(textFileName));

	Files.copy(TestJobInfo.JOB_LIB_JAR_PATH, userDirHasNotEntryClass.toPath().resolve(TestJobInfo.JOB_LIB_JAR_PATH.toFile().getName()));
	Files.createFile(userDirHasNotEntryClass.toPath().resolve(textFileName));

	final Path workingDirectory = FileUtils.getCurrentWorkingDirectory();
	Arrays.asList(userJarPath, userLibJarPath)
		.stream()
		.map(path -> FileUtils.relativizePath(workingDirectory, path))
		.map(FunctionUtils.uncheckedFunction(FileUtils::toURL))
		.forEach(expectedURLs::add);
}
 

@Override
public Collection<PluginDescriptor> findPlugins() throws IOException {

	if (!Files.isDirectory(pluginsRootDir)) {
		throw new IOException("Plugins root directory [" + pluginsRootDir + "] does not exist!");
	}

	return Files.list(pluginsRootDir)
		.filter((Path path) -> Files.isDirectory(path))
		.map(FunctionUtils.uncheckedFunction(this::createPluginDescriptorForSubDirectory))
		.collect(Collectors.toList());
}
 

private Map<String, FileStatus> getAllFilesInProvidedLibDirs(final List<Path> providedLibDirs) {
	final Map<String, FileStatus> allFiles = new HashMap<>();
	checkNotNull(providedLibDirs).forEach(
		FunctionUtils.uncheckedConsumer(
			path -> {
				if (!fileSystem.exists(path) || !fileSystem.isDirectory(path)) {
					LOG.warn("Provided lib dir {} does not exist or is not a directory. Ignoring.", path);
				} else {
					final RemoteIterator<LocatedFileStatus> iterable = fileSystem.listFiles(path, true);
					while (iterable.hasNext()) {
						final LocatedFileStatus locatedFileStatus = iterable.next();

						final String name = path.getParent().toUri()
								.relativize(locatedFileStatus.getPath().toUri())
								.toString();

						final FileStatus prevMapping = allFiles.put(name, locatedFileStatus);
						if (prevMapping != null) {
							throw new IOException(
								"Two files with the same filename exist in the shared libs: " +
									prevMapping.getPath() + " - " + locatedFileStatus.getPath() +
									". Please deduplicate.");
						}
					}

					if (LOG.isDebugEnabled()) {
						LOG.debug("The following files were found in the shared lib dir: {}",
								allFiles.values().stream()
										.map(fileStatus -> fileStatus.getPath().toString())
										.collect(Collectors.joining(", ")));
					}
				}
			})
	);
	return Collections.unmodifiableMap(allFiles);
}
 

private CompletableFuture<Void> waitForTerminatingJobManager(JobID jobId, JobGraph jobGraph, FunctionWithException<JobGraph, CompletableFuture<Void>, ?> action) {
	final CompletableFuture<Void> jobManagerTerminationFuture = getJobTerminationFuture(jobId)
		.exceptionally((Throwable throwable) -> {
			throw new CompletionException(
				new DispatcherException(
					String.format("Termination of previous JobManager for job %s failed. Cannot submit job under the same job id.", jobId),
					throwable)); });

	return jobManagerTerminationFuture.thenComposeAsync(
		FunctionUtils.uncheckedFunction((ignored) -> {
			jobManagerTerminationFutures.remove(jobId);
			return action.apply(jobGraph);
		}),
		getMainThreadExecutor());
}
 

private CompletableFuture<Void> submitAddedJob(JobGraph jobGraph) {
	final DispatcherGateway dispatcherGateway = getDispatcherGatewayInternal();

	return dispatcherGateway
		.submitJob(jobGraph, RpcUtils.INF_TIMEOUT)
		.thenApply(FunctionUtils.nullFn())
		.exceptionally(this::filterOutDuplicateJobSubmissionException);
}
 

private CompletableFuture<Void> revokeJobMasterLeadership() {
	log.info("JobManager for job {} ({}) at {} was revoked leadership.",
		jobGraph.getName(), jobGraph.getJobID(), jobMasterService.getAddress());

	setNewLeaderGatewayFuture();

	return jobMasterService
		.suspend(new FlinkException("JobManager is no longer the leader."))
		.thenApply(FunctionUtils.nullFn());
}
 

protected AbstractUserClassPathJobGraphRetriever(@Nullable File jobDir) throws IOException {
	if (jobDir == null) {
		userClassPaths = Collections.emptyList();
	} else {
		final Path workingDirectory = FileUtils.getCurrentWorkingDirectory();
		final Collection<URL> relativeJarURLs = FileUtils.listFilesInDirectory(jobDir.toPath(), FileUtils::isJarFile)
			.stream()
			.map(path -> FileUtils.relativizePath(workingDirectory, path))
			.map(FunctionUtils.uncheckedFunction(FileUtils::toURL))
			.collect(Collectors.toList());
		this.userClassPaths = Collections.unmodifiableCollection(relativeJarURLs);
	}
}
 

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

@Test
public void testGetUserClassPath() throws IOException {
	final File testJobDir = temporaryFolder.newFolder("_test_job");
	final Collection<Path> testFiles = FileUtilsTest.prepareTestFiles(testJobDir.toPath());
	final Path currentWorkingDirectory = FileUtils.getCurrentWorkingDirectory();
	final TestJobGraphRetriever testJobGraphRetriever = new TestJobGraphRetriever(testJobDir);

	assertThat(testJobGraphRetriever.getUserClassPaths(), containsInAnyOrder(testFiles.stream()
		.map(file -> FileUtils.relativizePath(currentWorkingDirectory, file))
		.map(FunctionUtils.uncheckedFunction(FileUtils::toURL)).toArray()));
}
 

/**
 * Executes the given supplier with the main thread executor until completion, returns the result or a exception.
 * This method blocks until the execution is complete.
 */
public <U> U execute(@Nonnull SupplierWithException<U, Throwable> supplierWithException) {
	return CompletableFuture.supplyAsync(
		FunctionUtils.uncheckedSupplier(supplierWithException),
		mainThreadExecutor)
		.join();
}
 

private Collection<URL> discoverUserClassPaths(@Nullable File jobDir) throws IOException {
	if (jobDir == null) {
		return Collections.emptyList();
	}

	final Path workingDirectory = FileUtils.getCurrentWorkingDirectory();
	final Collection<URL> relativeJarURLs = FileUtils.listFilesInDirectory(jobDir.toPath(), FileUtils::isJarFile)
			.stream()
			.map(path -> FileUtils.relativizePath(workingDirectory, path))
			.map(FunctionUtils.uncheckedFunction(FileUtils::toURL))
			.collect(Collectors.toList());
	return Collections.unmodifiableCollection(relativeJarURLs);
}
 

/**
 * Executes the given supplier with the main thread executor until completion, returns the result or a exception.
 * This method blocks until the execution is complete.
 */
public <U> U execute(@Nonnull SupplierWithException<U, Throwable> supplierWithException) {
	return CompletableFuture.supplyAsync(
		FunctionUtils.uncheckedSupplier(supplierWithException),
		mainThreadExecutor)
		.join();
}
 

@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<Void> waitForTerminatingJobManager(JobID jobId, JobGraph jobGraph, FunctionWithException<JobGraph, CompletableFuture<Void>, ?> action) {
	final CompletableFuture<Void> jobManagerTerminationFuture = getJobTerminationFuture(jobId)
		.exceptionally((Throwable throwable) -> {
			throw new CompletionException(
				new DispatcherException(
					String.format("Termination of previous JobManager for job %s failed. Cannot submit job under the same job id.", jobId),
					throwable)); });

	return jobManagerTerminationFuture.thenComposeAsync(
		FunctionUtils.uncheckedFunction((ignored) -> {
			jobManagerTerminationFutures.remove(jobId);
			return action.apply(jobGraph);
		}),
		getMainThreadExecutor());
}
 

private CompletableFuture<Void> revokeJobMasterLeadership() {
	log.info("JobManager for job {} ({}) was revoked leadership at {}.",
		jobGraph.getName(), jobGraph.getJobID(), getAddress());

	setNewLeaderGatewayFuture();

	return jobMasterService
		.suspend(new FlinkException("JobManager is no longer the leader."))
		.thenApply(FunctionUtils.nullFn());
}
 

private CompletableFuture<Void> revokeJobMasterLeadership() {
	log.info("JobManager for job {} ({}) was revoked leadership at {}.",
		jobGraph.getName(), jobGraph.getJobID(), getAddress());

	setNewLeaderGatewayFuture();

	return jobMasterService
		.suspend(new FlinkException("JobManager is no longer the leader."))
		.thenApply(FunctionUtils.nullFn());
}
 

private CompletableFuture<Void> waitForTerminatingJobManager(JobID jobId, JobGraph jobGraph, FunctionWithException<JobGraph, CompletableFuture<Void>, ?> action) {
	final CompletableFuture<Void> jobManagerTerminationFuture = getJobTerminationFuture(jobId)
		.exceptionally((Throwable throwable) -> {
			throw new CompletionException(
				new DispatcherException(
					String.format("Termination of previous JobManager for job %s failed. Cannot submit job under the same job id.", jobId),
					throwable)); });

	return jobManagerTerminationFuture.thenComposeAsync(
		FunctionUtils.uncheckedFunction((ignored) -> {
			jobManagerTerminationFutures.remove(jobId);
			return action.apply(jobGraph);
		}),
		getMainThreadExecutor());
}
 

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

@Override
public Collection<PluginDescriptor> findPlugins() throws IOException {

	if (!Files.isDirectory(pluginsRootDir)) {
		throw new IOException("Plugins root directory [" + pluginsRootDir + "] does not exist!");
	}

	return Files.list(pluginsRootDir)
		.filter((Path path) -> Files.isDirectory(path))
		.map(FunctionUtils.uncheckedFunction(this::createPluginDescriptorForSubDirectory))
		.collect(Collectors.toList());
}
 

/**
 * Executes the given supplier with the main thread executor until completion, returns the result or a exception.
 * This method blocks until the execution is complete.
 */
public <U> U execute(@Nonnull SupplierWithException<U, Throwable> supplierWithException) {
	return CompletableFuture.supplyAsync(
		FunctionUtils.uncheckedSupplier(supplierWithException),
		mainThreadExecutor)
		.join();
}
 

public CompletableFuture<URI> getRestAddress() {
	synchronized (lock) {
		checkState(running, "MiniCluster is not yet running.");
		return webMonitorLeaderRetriever.getLeaderFuture().thenApply(FunctionUtils.uncheckedFunction(addressLeaderIdTuple -> new URI(addressLeaderIdTuple.f0)));
	}
}
 

/**
 * Tests that pending slot requests are rejected if a slot report with a different allocation
 * is received.
 */
@Test
public void testSlotReportWhileActiveSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().build();

	final JobID jobId = new JobID();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile = ResourceProfile.fromResources(42.0, 1337);
	final SlotRequest slotRequest = new SlotRequest(jobId, allocationId, resourceProfile, "foobar");
	final CompletableFuture<Acknowledge> slotRequestFuture1 = new CompletableFuture<>();

	final Iterator<CompletableFuture<Acknowledge>> slotRequestFutureIterator = Arrays.asList(
		slotRequestFuture1,
		CompletableFuture.completedFuture(Acknowledge.get())).iterator();
	final ArrayBlockingQueue<SlotID> slotIds = new ArrayBlockingQueue<>(2);

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

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

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

	final ScheduledExecutor mainThreadExecutor = TestingUtils.defaultScheduledExecutor();

	final SlotManagerImpl slotManager = createSlotManagerBuilder()
		.setScheduledExecutor(mainThreadExecutor)
		.build();

	try {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		CompletableFuture<Void> registrationFuture = CompletableFuture.supplyAsync(
			() -> {
				slotManager.registerTaskManager(taskManagerConnection, slotReport);

				return null;
			},
			mainThreadExecutor)
		.thenAccept(
			(Object value) -> {
				try {
					slotManager.registerSlotRequest(slotRequest);
				} catch (ResourceManagerException e) {
					throw new RuntimeException("Could not register slots.", e);
				}
			});

		// check that no exception has been thrown
		registrationFuture.get();

		final SlotID requestedSlotId = slotIds.take();
		final SlotID freeSlotId = requestedSlotId.equals(slotId1) ? slotId2 : slotId1;

		final SlotStatus newSlotStatus1 = new SlotStatus(requestedSlotId, resourceProfile, new JobID(), new AllocationID());
		final SlotStatus newSlotStatus2 = new SlotStatus(freeSlotId, resourceProfile);
		final SlotReport newSlotReport = new SlotReport(Arrays.asList(newSlotStatus1, newSlotStatus2));

		CompletableFuture<Boolean> reportSlotStatusFuture = CompletableFuture.supplyAsync(
			// this should update the slot with the pending slot request triggering the reassignment of it
			() -> slotManager.reportSlotStatus(taskManagerConnection.getInstanceID(), newSlotReport),
			mainThreadExecutor);

		assertTrue(reportSlotStatusFuture.get());

		final SlotID requestedSlotId2 = slotIds.take();

		assertEquals(freeSlotId, requestedSlotId2);
	} finally {
		CompletableFuture.runAsync(
			ThrowingRunnable.unchecked(slotManager::close),
			mainThreadExecutor);
	}
}
 

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

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

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

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

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

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

		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		slotManager.registerSlotRequest(slotRequest);

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

		TaskManagerSlot failedSlot = slotManager.getSlot(firstSlotId);

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

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

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

		TaskManagerSlot slot = slotManager.getSlot(secondSlotId);

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

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

@Override
public CompletableFuture<Void> getTerminationFuture() {
	return executionGraph.getTerminationFuture().thenApply(FunctionUtils.nullFn());
}