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

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

/**
 * @return wrapped {@link OptionalFailure} returned by {@code valueSupplier} or wrapped failure if
 * {@code valueSupplier} has thrown an {@link Exception}.
 */
public static <T> OptionalFailure<T> createFrom(CheckedSupplier<T> valueSupplier) {
	try {
		return of(valueSupplier.get());
	} catch (Exception ex) {
		return ofFailure(ex);
	}
}
 

private <C> CompletableFuture<C> retry(
		CheckedSupplier<CompletableFuture<C>> operation,
		Predicate<Throwable> retryPredicate) {
	return FutureUtils.retryWithDelay(
		CheckedSupplier.unchecked(operation),
		restClusterClientConfiguration.getRetryMaxAttempts(),
		Time.milliseconds(restClusterClientConfiguration.getRetryDelay()),
		retryPredicate,
		new ScheduledExecutorServiceAdapter(retryExecutorService));
}
 

/**
 * Tests that the {@link SlotPoolImpl} completes stashed slot requests in request order.
 */
@Test
public void testStashOrderMaintainsRequestOrder() {
	runSlotRequestCompletionTest(
		CheckedSupplier.unchecked(this::setUpSlotPool),
		this::connectToResourceManager);
}
 

/**
 * Tests that the {@link SlotPoolImpl} completes slots in request order.
 */
@Test
public void testRequestsAreCompletedInRequestOrder() {
	runSlotRequestCompletionTest(
		CheckedSupplier.unchecked(this::setUpSlotPoolAndConnectToResourceManager),
		slotPool -> {});
}
 

/**
 * Tests that we can concurrently create two {@link ActorSystem} without port conflicts.
 * This effectively tests that we don't open a socket to check for a ports availability.
 * See FLINK-10580 for more details.
 */
@Test
public void testConcurrentActorSystemCreation() throws Exception {
	final int concurrentCreations = 10;
	final ExecutorService executorService = Executors.newFixedThreadPool(concurrentCreations);
	final CyclicBarrier cyclicBarrier = new CyclicBarrier(concurrentCreations);

	try {
		final List<CompletableFuture<Void>> actorSystemFutures = IntStream.range(0, concurrentCreations)
			.mapToObj(
				ignored ->
					CompletableFuture.supplyAsync(
						CheckedSupplier.unchecked(() -> {
							cyclicBarrier.await();

							return BootstrapTools.startRemoteActorSystem(
								new Configuration(),
								"localhost",
								"0",
								LOG);
						}), executorService))
			.map(
				// terminate ActorSystems
				actorSystemFuture ->
					actorSystemFuture.thenCompose(AkkaUtils::terminateActorSystem)
			).collect(Collectors.toList());

		FutureUtils.completeAll(actorSystemFutures).get();
	} finally {
		ExecutorUtils.gracefulShutdown(10000L, TimeUnit.MILLISECONDS, executorService);
	}
}
 

private Map<StateHandleID, CompletableFuture<StreamStateHandle>> createUploadFutures(
	Map<StateHandleID, Path> files,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) {
	Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = new HashMap<>(files.size());

	for (Map.Entry<StateHandleID, Path> entry : files.entrySet()) {
		final Supplier<StreamStateHandle> supplier =
			CheckedSupplier.unchecked(() -> uploadLocalFileToCheckpointFs(entry.getValue(), checkpointStreamFactory, closeableRegistry));
		futures.put(entry.getKey(), CompletableFuture.supplyAsync(supplier, executorService));
	}

	return futures;
}
 

/**
 * @return wrapped {@link OptionalFailure} returned by {@code valueSupplier} or wrapped failure if
 * {@code valueSupplier} has thrown an {@link Exception}.
 */
public static <T> OptionalFailure<T> createFrom(CheckedSupplier<T> valueSupplier) {
	try {
		return of(valueSupplier.get());
	} catch (Exception ex) {
		return ofFailure(ex);
	}
}
 

private Collection<JobID> waitForRunningJobs(ClusterClient<?> clusterClient, Time timeout) throws ExecutionException, InterruptedException {
	return FutureUtils.retrySuccessfulWithDelay(
			CheckedSupplier.unchecked(clusterClient::listJobs),
			Time.milliseconds(50L),
			Deadline.fromNow(Duration.ofMillis(timeout.toMilliseconds())),
			jobs -> !jobs.isEmpty(),
			TestingUtils.defaultScheduledExecutor())
		.get()
		.stream()
		.map(JobStatusMessage::getJobId)
		.collect(Collectors.toList());
}
 

private <C> CompletableFuture<C> retry(
		CheckedSupplier<CompletableFuture<C>> operation,
		Predicate<Throwable> retryPredicate) {
	return FutureUtils.retryWithDelay(
		CheckedSupplier.unchecked(operation),
		restClusterClientConfiguration.getRetryMaxAttempts(),
		Time.milliseconds(restClusterClientConfiguration.getRetryDelay()),
		retryPredicate,
		new ScheduledExecutorServiceAdapter(retryExecutorService));
}
 

/**
 * Tests that the {@link SlotPoolImpl} completes stashed slot requests in request order.
 */
@Test
public void testStashOrderMaintainsRequestOrder() throws Exception {
	runSlotRequestCompletionTest(
		CheckedSupplier.unchecked(this::setUpSlotPool),
		this::connectToResourceManager);
}
 

/**
 * Tests that the {@link SlotPoolImpl} completes slots in request order.
 */
@Test
public void testRequestsAreCompletedInRequestOrder() throws Exception {
	runSlotRequestCompletionTest(
		CheckedSupplier.unchecked(this::setUpSlotPoolAndConnectToResourceManager),
		slotPool -> {});
}
 

private Map<StateHandleID, CompletableFuture<StreamStateHandle>> createUploadFutures(
	Map<StateHandleID, Path> files,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) {
	Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = new HashMap<>(files.size());

	for (Map.Entry<StateHandleID, Path> entry : files.entrySet()) {
		final Supplier<StreamStateHandle> supplier =
			CheckedSupplier.unchecked(() -> uploadLocalFileToCheckpointFs(entry.getValue(), checkpointStreamFactory, closeableRegistry));
		futures.put(entry.getKey(), CompletableFuture.supplyAsync(supplier, executorService));
	}

	return futures;
}
 

private Collection<JobID> waitForRunningJobs(ClusterClient<?> clusterClient, Time timeout) throws ExecutionException, InterruptedException {
	return FutureUtils.retrySuccessfulWithDelay(
			CheckedSupplier.unchecked(clusterClient::listJobs),
			Time.milliseconds(50L),
			Deadline.fromNow(Duration.ofMillis(timeout.toMilliseconds())),
			jobs -> !jobs.isEmpty(),
			TestingUtils.defaultScheduledExecutor())
		.get()
		.stream()
		.map(JobStatusMessage::getJobId)
		.collect(Collectors.toList());
}
 

private Map<StateHandleID, CompletableFuture<StreamStateHandle>> createUploadFutures(
	Map<StateHandleID, Path> files,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) {
	Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures = new HashMap<>(files.size());

	for (Map.Entry<StateHandleID, Path> entry : files.entrySet()) {
		final Supplier<StreamStateHandle> supplier =
			CheckedSupplier.unchecked(() -> uploadLocalFileToCheckpointFs(entry.getValue(), checkpointStreamFactory, closeableRegistry));
		futures.put(entry.getKey(), CompletableFuture.supplyAsync(supplier, executorService));
	}

	return futures;
}
 

/**
 * Tests that we can concurrently create two {@link ActorSystem} without port conflicts.
 * This effectively tests that we don't open a socket to check for a ports availability.
 * See FLINK-10580 for more details.
 */
@Test
public void testConcurrentActorSystemCreation() throws Exception {
	final int concurrentCreations = 10;
	final ExecutorService executorService = Executors.newFixedThreadPool(concurrentCreations);
	final CyclicBarrier cyclicBarrier = new CyclicBarrier(concurrentCreations);

	try {
		final List<CompletableFuture<Void>> actorSystemFutures = IntStream.range(0, concurrentCreations)
			.mapToObj(
				ignored ->
					CompletableFuture.supplyAsync(
						CheckedSupplier.unchecked(() -> {
							cyclicBarrier.await();

							return BootstrapTools.startActorSystem(
								new Configuration(),
								"localhost",
								"0",
								LOG);
						}), executorService))
			.map(
				// terminate ActorSystems
				actorSystemFuture ->
					actorSystemFuture.thenCompose(AkkaUtils::terminateActorSystem)
			).collect(Collectors.toList());

		FutureUtils.completeAll(actorSystemFutures).get();
	} finally {
		ExecutorUtils.gracefulShutdown(10000L, TimeUnit.MILLISECONDS, executorService);
	}
}
 

private <C> CompletableFuture<C> retry(
		CheckedSupplier<CompletableFuture<C>> operation,
		Predicate<Throwable> retryPredicate) {
	return FutureUtils.retryWithDelay(
		CheckedSupplier.unchecked(operation),
		restClusterClientConfiguration.getRetryMaxAttempts(),
		Time.milliseconds(restClusterClientConfiguration.getRetryDelay()),
		retryPredicate,
		new ScheduledExecutorServiceAdapter(retryExecutorService));
}
 

private Collection<JobID> waitForRunningJobs(ClusterClient<?> clusterClient, Time timeout) throws ExecutionException, InterruptedException {
	return FutureUtils.retrySuccessfulWithDelay(
			CheckedSupplier.unchecked(clusterClient::listJobs),
			Time.milliseconds(50L),
			Deadline.fromNow(Duration.ofMillis(timeout.toMilliseconds())),
			jobs -> !jobs.isEmpty(),
			TestingUtils.defaultScheduledExecutor())
		.get()
		.stream()
		.map(JobStatusMessage::getJobId)
		.collect(Collectors.toList());
}
 

/**
 * @return wrapped {@link OptionalFailure} returned by {@code valueSupplier} or wrapped failure if
 * {@code valueSupplier} has thrown an {@link Exception}.
 */
public static <T> OptionalFailure<T> createFrom(CheckedSupplier<T> valueSupplier) {
	try {
		return of(valueSupplier.get());
	} catch (Exception ex) {
		return ofFailure(ex);
	}
}
 

/**
 * Tests that we can concurrently create two {@link ActorSystem} without port conflicts.
 * This effectively tests that we don't open a socket to check for a ports availability.
 * See FLINK-10580 for more details.
 */
@Test
public void testConcurrentActorSystemCreation() throws Exception {
	final int concurrentCreations = 10;
	final ExecutorService executorService = Executors.newFixedThreadPool(concurrentCreations);
	final CyclicBarrier cyclicBarrier = new CyclicBarrier(concurrentCreations);

	try {
		final List<CompletableFuture<Void>> actorSystemFutures = IntStream.range(0, concurrentCreations)
			.mapToObj(
				ignored ->
					CompletableFuture.supplyAsync(
						CheckedSupplier.unchecked(() -> {
							cyclicBarrier.await();

							return BootstrapTools.startActorSystem(
								new Configuration(),
								"localhost",
								"0",
								LOG);
						}), executorService))
			.map(
				// terminate ActorSystems
				actorSystemFuture ->
					actorSystemFuture.thenCompose(AkkaUtils::terminateActorSystem)
			).collect(Collectors.toList());

		FutureUtils.completeAll(actorSystemFutures).get();
	} finally {
		ExecutorUtils.gracefulShutdown(10000L, TimeUnit.MILLISECONDS, executorService);
	}
}
 

/**
 * Tests that a batch job can be executed with fewer slots than its parallelism.
 * See FLINK-13187 for more information.
 */
@Test
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception {
	final int parallelism = 5;
	final Time batchSlotTimeout = Time.milliseconds(5L);
	final JobGraph jobGraph = createJobGraph(parallelism);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST);

	try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) {
		final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism);
		TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setSubmitTaskConsumer(
				(tdd, ignored) -> {
					submittedTasksQueue.offer(tdd.getExecutionAttemptId());
					return CompletableFuture.completedFuture(Acknowledge.get());
				})
			.createTestingTaskExecutorGateway();

		// register a single slot at the slot pool
		SlotPoolUtils.offerSlots(
			slotPool,
			mainThreadExecutor,
			Collections.singletonList(ResourceProfile.ANY),
			new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate()));

		final LegacyScheduler legacyScheduler = createLegacyScheduler(jobGraph, slotPool, mainThreadExecutor, batchSlotTimeout);

		final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener();
		legacyScheduler.registerJobStatusListener(jobStatusListener);
		startScheduling(legacyScheduler, mainThreadExecutor);

		// wait until the batch slot timeout has been reached
		Thread.sleep(batchSlotTimeout.toMilliseconds());

		final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture();

		for (int i = 0; i < parallelism; i++) {
			final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take));

			// wait until one of them is completed
			CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join();

			if (submittedTaskFuture.isDone()) {
				finishExecution(submittedTaskFuture.get(), legacyScheduler, mainThreadExecutor);
			} else {
				fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get()));
			}
		}

		assertThat(terminationFuture.get(), is(JobStatus.FINISHED));
	}
}
 

/**
 * Tests that we fail the operation if the client closes.
 */
@Test
public void testRestClientClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);

	Socket connectionSocket = null;

	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		restClient.close();

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	} finally {
		if (connectionSocket != null) {
			connectionSocket.close();
		}
	}
}
 

/**
 * Tests that we fail the operation if the remote connection closes.
 */
@Test
public void testConnectionClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);
	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		Socket connectionSocket = null;

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		if (connectionSocket != null) {
			// close connection
			connectionSocket.close();
		}

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	}
}
 

/**
 * Tests that we fail the operation if the client closes.
 */
@Test
public void testRestClientClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);

	Socket connectionSocket = null;

	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		restClient.close();

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	} finally {
		if (connectionSocket != null) {
			connectionSocket.close();
		}
	}
}
 

/**
 * Tests that we fail the operation if the remote connection closes.
 */
@Test
public void testConnectionClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);
	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		Socket connectionSocket = null;

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		if (connectionSocket != null) {
			// close connection
			connectionSocket.close();
		}

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	}
}
 

/**
 * Tests that we fail the operation if the remote connection closes.
 */
@Test
public void testConnectionClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);
	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		Socket connectionSocket = null;

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		if (connectionSocket != null) {
			// close connection
			connectionSocket.close();
		}

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	}
}
 

/**
 * Tests that we fail the operation if the client closes.
 */
@Test
public void testRestClientClosedHandling() throws Exception {
	final Configuration config = new Configuration();
	config.setLong(RestOptions.IDLENESS_TIMEOUT, 5000L);

	Socket connectionSocket = null;

	try (final ServerSocket serverSocket = new ServerSocket(0);
		final RestClient restClient = new RestClient(RestClientConfiguration.fromConfiguration(config), TestingUtils.defaultExecutor())) {

		final String targetAddress = "localhost";
		final int targetPort = serverSocket.getLocalPort();

		// start server
		final CompletableFuture<Socket> socketCompletableFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(serverSocket::accept));

		final CompletableFuture<EmptyResponseBody> responseFuture = restClient.sendRequest(
			targetAddress,
			targetPort,
			new TestMessageHeaders(),
			EmptyMessageParameters.getInstance(),
			EmptyRequestBody.getInstance(),
			Collections.emptyList());

		try {
			connectionSocket = socketCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
		} catch (TimeoutException ignored) {
			// could not establish a server connection --> see that the response failed
			socketCompletableFuture.cancel(true);
		}

		restClient.close();

		try {
			responseFuture.get();
		} catch (ExecutionException ee) {
			if (!ExceptionUtils.findThrowable(ee, IOException.class).isPresent()) {
				throw ee;
			}
		}
	} finally {
		if (connectionSocket != null) {
			connectionSocket.close();
		}
	}
}
 

/**
 * Tests that a batch job can be executed with fewer slots than its parallelism.
 * See FLINK-13187 for more information.
 */
@Test
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception {
	final int parallelism = 5;
	final Time batchSlotTimeout = Time.milliseconds(5L);
	final JobGraph jobGraph = createJobGraph(parallelism);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST);

	try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) {
		final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism);
		TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setSubmitTaskConsumer(
				(tdd, ignored) -> {
					submittedTasksQueue.offer(tdd.getExecutionAttemptId());
					return CompletableFuture.completedFuture(Acknowledge.get());
				})
			.createTestingTaskExecutorGateway();

		// register a single slot at the slot pool
		SlotPoolUtils.offerSlots(
			slotPool,
			mainThreadExecutor,
			Collections.singletonList(ResourceProfile.ANY),
			new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate()));

		final SlotProvider slotProvider = createSlotProvider(slotPool, mainThreadExecutor);
		final SchedulerNG scheduler = createScheduler(jobGraph, slotProvider, batchSlotTimeout);

		final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener();
		scheduler.registerJobStatusListener(jobStatusListener);
		startScheduling(scheduler, mainThreadExecutor);

		// wait until the batch slot timeout has been reached
		Thread.sleep(batchSlotTimeout.toMilliseconds());

		final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture();

		for (int i = 0; i < parallelism; i++) {
			final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take));

			// wait until one of them is completed
			CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join();

			if (submittedTaskFuture.isDone()) {
				finishExecution(submittedTaskFuture.get(), scheduler, mainThreadExecutor);
			} else {
				fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get()));
			}
		}

		assertThat(terminationFuture.get(), is(JobStatus.FINISHED));
	}
}
 

/**
 * Tests the consumption of multiple subpartitions via local input channels.
 *
 * <p>Multiple producer tasks produce pipelined partitions, which are consumed by multiple
 * tasks via local input channels.
 */
@Test
public void testConcurrentConsumeMultiplePartitions() throws Exception {
	// Config
	final int parallelism = 32;
	final int producerBufferPoolSize = parallelism + 1;
	final int numberOfBuffersPerChannel = 1024;

	checkArgument(parallelism >= 1);
	checkArgument(producerBufferPoolSize >= parallelism);
	checkArgument(numberOfBuffersPerChannel >= 1);

	// Setup
	// One thread per produced partition and one per consumer
	final ExecutorService executor = Executors.newFixedThreadPool(2 * parallelism);

	final NetworkBufferPool networkBuffers = new NetworkBufferPool(
		(parallelism * producerBufferPoolSize) + (parallelism * parallelism),
		TestBufferFactory.BUFFER_SIZE, 1);

	final ResultPartitionManager partitionManager = new ResultPartitionManager();

	final ResultPartitionID[] partitionIds = new ResultPartitionID[parallelism];
	final TestPartitionProducer[] partitionProducers = new TestPartitionProducer[parallelism];

	// Create all partitions
	for (int i = 0; i < parallelism; i++) {
		partitionIds[i] = new ResultPartitionID();

		final ResultPartition partition = new ResultPartitionBuilder()
			.setResultPartitionId(partitionIds[i])
			.setNumberOfSubpartitions(parallelism)
			.setNumTargetKeyGroups(parallelism)
			.setResultPartitionManager(partitionManager)
			.setBufferPoolFactory(p ->
				networkBuffers.createBufferPool(producerBufferPoolSize, producerBufferPoolSize))
			.build();

		// Create a buffer pool for this partition
		partition.setup();

		// Create the producer
		partitionProducers[i] = new TestPartitionProducer(
			partition,
			false,
			new TestPartitionProducerBufferSource(
				parallelism,
				partition.getBufferPool(),
				numberOfBuffersPerChannel)
		);
	}

	// Test
	try {
		// Submit producer tasks
		List<CompletableFuture<?>> results = Lists.newArrayListWithCapacity(
			parallelism + 1);

		for (int i = 0; i < parallelism; i++) {
			results.add(CompletableFuture.supplyAsync(
				CheckedSupplier.unchecked(partitionProducers[i]::call), executor));
		}

		// Submit consumer
		for (int i = 0; i < parallelism; i++) {
			final TestLocalInputChannelConsumer consumer = new TestLocalInputChannelConsumer(
				i,
				parallelism,
				numberOfBuffersPerChannel,
				networkBuffers.createBufferPool(parallelism, parallelism),
				partitionManager,
				new TaskEventDispatcher(),
				partitionIds);

			results.add(CompletableFuture.supplyAsync(CheckedSupplier.unchecked(consumer::call), executor));
		}

		FutureUtils.waitForAll(results).get();
	}
	finally {
		networkBuffers.destroyAllBufferPools();
		networkBuffers.destroy();
		executor.shutdown();
	}
}