org.apache.flink.client.ClientUtils Java Examples

/**
 * Override configuration settings by specified command line options.
 *
 * @param commandLine containing the overriding values
 * @return Effective configuration with the overridden configuration settings
 */
protected Configuration applyCommandLineOptionsToConfiguration(CommandLine commandLine) throws FlinkException {
	final Configuration resultingConfiguration = new Configuration(configuration);

	if (commandLine.hasOption(addressOption.getOpt())) {
		String addressWithPort = commandLine.getOptionValue(addressOption.getOpt());
		InetSocketAddress jobManagerAddress = ClientUtils.parseHostPortAddress(addressWithPort);
		setJobManagerAddressInConfig(resultingConfiguration, jobManagerAddress);
	}

	if (commandLine.hasOption(zookeeperNamespaceOption.getOpt())) {
		String zkNamespace = commandLine.getOptionValue(zookeeperNamespaceOption.getOpt());
		resultingConfiguration.setString(HighAvailabilityOptions.HA_CLUSTER_ID, zkNamespace);
	}

	return resultingConfiguration;
}
 

private void launchMultiExecuteJob(final boolean enforceSingleJobExecution) throws ProgramInvocationException {
	try (final ClusterClient<?> clusterClient =
				new MiniClusterClient(new Configuration(), MINI_CLUSTER_RESOURCE.getMiniCluster())) {

		final PackagedProgram program = PackagedProgram.newBuilder()
				.setEntryPointClassName(TestMultiExecute.class.getName())
				.build();

		final Configuration configuration = fromPackagedProgram(program, 1, false);

		ClientUtils.executeProgram(
				new TestExecutorServiceLoader(clusterClient, plan),
				configuration,
				program,
				enforceSingleJobExecution,
				false);
	}
}
 

private List<JobID> tryExecuteJobs(final DispatcherGateway dispatcherGateway, final PackagedProgram program, final Configuration configuration) {
	configuration.set(DeploymentOptions.ATTACHED, false);

	final List<JobID> applicationJobIds = new ArrayList<>();
	final PipelineExecutorServiceLoader executorServiceLoader =
			new WebSubmissionExecutorServiceLoader(applicationJobIds, dispatcherGateway);

	try {
		ClientUtils.executeProgram(executorServiceLoader, configuration, program, enforceSingleJobExecution, true);
	} catch (ProgramInvocationException e) {
		LOG.warn("Could not execute application: ", e);
		throw new FlinkRuntimeException("Could not execute application.", e);
	}

	return applicationJobIds;
}
 

public <T> String takeSavepoint(Collection<T> data, Function<SourceFunction<T>, StreamExecutionEnvironment> jobGraphFactory) throws Exception {

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableClosureCleaner();

		WaitingSource<T> waitingSource = createSource(data);

		JobGraph jobGraph = jobGraphFactory.apply(waitingSource).getStreamGraph().getJobGraph();
		JobID jobId = jobGraph.getJobID();

		ClusterClient<?> client = miniClusterResource.getClusterClient();

		try {
			JobSubmissionResult result = ClientUtils.submitJob(client, jobGraph);

			return CompletableFuture
				.runAsync(waitingSource::awaitSource)
				.thenCompose(ignore -> triggerSavepoint(client, result.getJobID()))
				.get(5, TimeUnit.MINUTES);
		} catch (Exception e) {
			throw new RuntimeException("Failed to take savepoint", e);
		} finally {
			client.cancel(jobId);
		}
	}
 

/**
 * This test verifies that the local execution environment cannot be created when
 * the program is submitted through a client.
 */
@Test
public void tryLocalExecution() throws ProgramInvocationException, ProgramMissingJobException {
	PackagedProgram packagedProgramMock = mock(PackagedProgram.class);

	when(packagedProgramMock.getUserCodeClassLoader())
			.thenReturn(packagedProgramMock.getClass().getClassLoader());

	doAnswer(new Answer<Void>() {
		@Override
		public Void answer(InvocationOnMock invocation) throws Throwable {
			ExecutionEnvironment.createLocalEnvironment();
			return null;
		}
	}).when(packagedProgramMock).invokeInteractiveModeForExecution();

	try {
		final ClusterClient<?> client = new MiniClusterClient(new Configuration(), MINI_CLUSTER_RESOURCE.getMiniCluster());
		final Configuration configuration = fromPackagedProgram(packagedProgramMock, 1, true);
		ClientUtils.executeProgram(new TestExecutorServiceLoader(client, plan), configuration, packagedProgramMock, false, false);
		fail("Creating the local execution environment should not be possible");
	}
	catch (InvalidProgramException e) {
		// that is what we want
	}
}
 

/**
 * Tests that we can submit a jobGraph in detached mode.
 */
@Test
public void testDetachedJobSubmission() throws Exception {

	final TestJobSubmitHandler testJobSubmitHandler = new TestJobSubmitHandler();

	try (TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(
		testJobSubmitHandler)) {
		RestClusterClient<?> restClusterClient = createRestClusterClient(restServerEndpoint.getServerAddress().getPort());

		try {
			final JobSubmissionResult jobSubmissionResult = ClientUtils.submitJob(restClusterClient, jobGraph);

			// if the detached mode didn't work, then we would not reach this point because the execution result
			// retrieval would have failed.
			assertThat(jobSubmissionResult, is(instanceOf(DetachedJobExecutionResult.class)));
			assertThat(jobSubmissionResult.getJobID(), is(jobId));
		} finally {
			restClusterClient.close();
		}
	}

}
 

private String submitJobAndTakeSavepoint(MiniClusterResourceFactory clusterFactory, int parallelism) throws Exception {
	final JobGraph jobGraph = createJobGraph(parallelism, 0, 1000);
	final JobID jobId = jobGraph.getJobID();
	StatefulCounter.resetForTest(parallelism);

	MiniClusterWithClientResource cluster = clusterFactory.get();
	cluster.before();
	ClusterClient<?> client = cluster.getClusterClient();

	try {
		ClientUtils.submitJob(client, jobGraph);

		StatefulCounter.getProgressLatch().await();

		return client.cancelWithSavepoint(jobId, null).get();
	} finally {
		cluster.after();
		StatefulCounter.resetForTest(parallelism);
	}
}
 

private void restoreJob(ClusterClient<?> clusterClient, Deadline deadline, String savepointPath) throws Exception {
	JobGraph jobToRestore = createJobGraph(ExecutionMode.RESTORE);
	jobToRestore.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepointPath, allowNonRestoredState));

	assertNotNull("Job doesn't have a JobID.", jobToRestore.getJobID());

	ClientUtils.submitJob(clusterClient, jobToRestore);

	CompletableFuture<JobStatus> jobStatusFuture = FutureUtils.retrySuccessfulWithDelay(
		() -> clusterClient.getJobStatus(jobToRestore.getJobID()),
		Time.milliseconds(50),
		deadline,
		(jobStatus) -> jobStatus == JobStatus.FINISHED,
		TestingUtils.defaultScheduledExecutor());
	assertEquals(
		JobStatus.FINISHED,
		jobStatusFuture.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS));
}
 

/**
 * Override configuration settings by specified command line options.
 *
 * @param commandLine containing the overriding values
 * @return Effective configuration with the overridden configuration settings
 */
protected Configuration applyCommandLineOptionsToConfiguration(CommandLine commandLine) throws FlinkException {
	final Configuration resultingConfiguration = new Configuration(configuration);

	if (commandLine.hasOption(addressOption.getOpt())) {
		String addressWithPort = commandLine.getOptionValue(addressOption.getOpt());
		InetSocketAddress jobManagerAddress = ClientUtils.parseHostPortAddress(addressWithPort);
		setJobManagerAddressInConfig(resultingConfiguration, jobManagerAddress);
	}

	if (commandLine.hasOption(zookeeperNamespaceOption.getOpt())) {
		String zkNamespace = commandLine.getOptionValue(zookeeperNamespaceOption.getOpt());
		resultingConfiguration.setString(HighAvailabilityOptions.HA_CLUSTER_ID, zkNamespace);
	}

	return resultingConfiguration;
}
 

/**
 * Tests simple value state queryable state instance. Each source emits
 * (subtaskIndex, 0)..(subtaskIndex, numElements) tuples, which are then
 * queried. The tests succeeds after each subtask index is queried with
 * value numElements (the latest element updated the state).
 */
@Test
public void testValueState() throws Exception {
	final Deadline deadline = Deadline.now().plus(TEST_TIMEOUT);
	final long numElements = 1024L;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setStateBackend(stateBackend);
	env.setParallelism(maxParallelism);
	// Very important, because cluster is shared between tests and we
	// don't explicitly check that all slots are available before
	// submitting.
	env.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000L));

	DataStream<Tuple2<Integer, Long>> source = env.addSource(new TestAscendingValueSource(numElements));

	// Value state
	ValueStateDescriptor<Tuple2<Integer, Long>> valueState = new ValueStateDescriptor<>("any", source.getType());

	source.keyBy(new KeySelector<Tuple2<Integer, Long>, Integer>() {
		private static final long serialVersionUID = 7662520075515707428L;

		@Override
		public Integer getKey(Tuple2<Integer, Long> value) {
			return value.f0;
		}
	}).asQueryableState("hakuna", valueState);

	try (AutoCancellableJob autoCancellableJob = new AutoCancellableJob(deadline, clusterClient, env)) {

		final JobID jobId = autoCancellableJob.getJobId();
		final JobGraph jobGraph = autoCancellableJob.getJobGraph();

		ClientUtils.submitJob(clusterClient, jobGraph);
		executeValueQuery(deadline, client, jobId, "hakuna", valueState, numElements);
	}
}
 

/**
 * Runs the user program entrypoint and completes the given {@code jobIdsFuture} with the {@link
 * JobID JobIDs} of the submitted jobs.
 *
 * <p>This should be executed in a separate thread (or task).
 */
private void runApplicationEntryPoint(
		final CompletableFuture<List<JobID>> jobIdsFuture,
		final DispatcherGateway dispatcher,
		final ScheduledExecutor scheduledExecutor,
		final boolean enforceSingleJobExecution) {
	try {
		final List<JobID> applicationJobIds =
				new ArrayList<>(getRecoveredJobIds(recoveredJobs));

		final PipelineExecutorServiceLoader executorServiceLoader =
				new EmbeddedExecutorServiceLoader(
						applicationJobIds, dispatcher, scheduledExecutor);

		ClientUtils.executeProgram(
				executorServiceLoader,
				configuration,
				application,
				enforceSingleJobExecution,
				true /* suppress sysout */);

		if (applicationJobIds.isEmpty()) {
			jobIdsFuture.completeExceptionally(
					new ApplicationExecutionException(
							"The application contains no execute() calls."));
		} else {
			jobIdsFuture.complete(applicationJobIds);
		}
	} catch (Throwable t) {
		jobIdsFuture.completeExceptionally(
				new ApplicationExecutionException("Could not execute application.", t));
	}
}
 

/**
 * This test verifies correct job submission messaging logic and plan translation calls.
 */
@Test
public void shouldSubmitToJobClient() throws Exception {
	final ClusterClient<?> clusterClient = new MiniClusterClient(new Configuration(), MINI_CLUSTER_RESOURCE.getMiniCluster());
	JobGraph jobGraph = FlinkPipelineTranslationUtil.getJobGraph(
			plan,
			new Configuration(),
			1);

	jobGraph.addJars(Collections.emptyList());
	jobGraph.setClasspaths(Collections.emptyList());

	JobSubmissionResult result = ClientUtils.submitJob(clusterClient, jobGraph);
	assertNotNull(result);
}
 

private void validateModification(String savepointPath) throws ProgramInvocationException {
	StreamExecutionEnvironment sEnv = StreamExecutionEnvironment.getExecutionEnvironment();
	sEnv.setStateBackend(backend);

	CollectSink.accountList.clear();

	DataStream<Account> stream = sEnv.fromCollection(accounts)
		.keyBy(acc -> acc.id)
		.flatMap(new UpdateAndGetAccount())
		.uid(ACCOUNT_UID);

	stream.addSink(new CollectSink());

	stream
		.map(acc -> acc.id)
		.map(new StatefulOperator())
		.uid(MODIFY_UID)
		.addSink(new DiscardingSink<>());

	JobGraph jobGraph = sEnv.getStreamGraph().getJobGraph();
	jobGraph.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepointPath, false));

	ClusterClient<?> client = miniClusterResource.getClusterClient();
	ClientUtils.submitJobAndWaitForResult(client, jobGraph, SavepointWriterITCase.class.getClassLoader());

	Assert.assertEquals("Unexpected output", 3, CollectSink.accountList.size());
}
 

private void validateBootstrap(String savepointPath) throws ProgramInvocationException {
	StreamExecutionEnvironment sEnv = StreamExecutionEnvironment.getExecutionEnvironment();
	sEnv.setStateBackend(backend);

	CollectSink.accountList.clear();

	sEnv.fromCollection(accounts)
		.keyBy(acc -> acc.id)
		.flatMap(new UpdateAndGetAccount())
		.uid(ACCOUNT_UID)
		.addSink(new CollectSink());

	sEnv
		.fromCollection(currencyRates)
		.connect(sEnv.fromCollection(currencyRates).broadcast(descriptor))
		.process(new CurrencyValidationFunction())
		.uid(CURRENCY_UID)
		.addSink(new DiscardingSink<>());

	JobGraph jobGraph = sEnv.getStreamGraph().getJobGraph();
	jobGraph.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepointPath, false));

	ClusterClient<?> client = miniClusterResource.getClusterClient();
	ClientUtils.submitJobAndWaitForResult(client, jobGraph, SavepointWriterITCase.class.getClassLoader());

	Assert.assertEquals("Unexpected output", 3, CollectSink.accountList.size());
}
 

/**
 * Tests simple value state queryable state instance. Each source emits
 * (subtaskIndex, 0)..(subtaskIndex, numElements) tuples, which are then
 * queried. The tests succeeds after each subtask index is queried with
 * value numElements (the latest element updated the state).
 *
 * <p>This is the same as the simple value state test, but uses the API shortcut.
 */
@Test
public void testValueStateShortcut() throws Exception {
	final Deadline deadline = Deadline.now().plus(TEST_TIMEOUT);
	final long numElements = 1024L;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setStateBackend(stateBackend);
	env.setParallelism(maxParallelism);
	// Very important, because cluster is shared between tests and we
	// don't explicitly check that all slots are available before
	// submitting.
	env.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000L));

	DataStream<Tuple2<Integer, Long>> source = env.addSource(new TestAscendingValueSource(numElements));

	// Value state shortcut
	final QueryableStateStream<Integer, Tuple2<Integer, Long>> queryableState =
			source.keyBy(new KeySelector<Tuple2<Integer, Long>, Integer>() {
				private static final long serialVersionUID = 9168901838808830068L;

				@Override
				public Integer getKey(Tuple2<Integer, Long> value) {
					return value.f0;
				}
			}).asQueryableState("matata");

	@SuppressWarnings("unchecked")
	final ValueStateDescriptor<Tuple2<Integer, Long>> stateDesc =
			(ValueStateDescriptor<Tuple2<Integer, Long>>) queryableState.getStateDescriptor();

	try (AutoCancellableJob autoCancellableJob = new AutoCancellableJob(deadline, clusterClient, env)) {

		final JobID jobId = autoCancellableJob.getJobId();
		final JobGraph jobGraph = autoCancellableJob.getJobGraph();

		ClientUtils.submitJob(clusterClient, jobGraph);
		executeValueQuery(deadline, client, jobId, "matata", stateDesc, numElements);
	}
}
 

private String takeSavepoint(JobGraph jobGraph) throws Exception {
	SavepointSource.initializeForTest();

	ClusterClient<?> client = miniClusterResource.getClusterClient();
	JobID jobId = jobGraph.getJobID();

	Deadline deadline = Deadline.fromNow(Duration.ofMinutes(5));

	String dirPath = getTempDirPath(new AbstractID().toHexString());

	try {
		JobSubmissionResult result = ClientUtils.submitJob(client, jobGraph);

		boolean finished = false;
		while (deadline.hasTimeLeft()) {
			if (SavepointSource.isFinished()) {
				finished = true;

				break;
			}

			try {
				Thread.sleep(2L);
			} catch (InterruptedException ignored) {
				Thread.currentThread().interrupt();
			}
		}

		if (!finished) {
			Assert.fail("Failed to initialize state within deadline");
		}

		CompletableFuture<String> path = client.triggerSavepoint(result.getJobID(), dirPath);
		return path.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
	} finally {
		client.cancel(jobId).get();
	}
}
 

@Test
public void testJobRetrieval() throws Exception {
	final JobID jobID = new JobID();

	final JobVertex imalock = new JobVertex("imalock");
	imalock.setInvokableClass(SemaphoreInvokable.class);

	final JobGraph jobGraph = new JobGraph(jobID, "testjob", imalock);

	// acquire the lock to make sure that the job cannot complete until the job client
	// has been attached in resumingThread
	lock.acquire();

	ClientUtils.submitJob(client, jobGraph);

	final CheckedThread resumingThread = new CheckedThread("Flink-Job-Retriever") {
		@Override
		public void go() throws Exception {
			assertNotNull(client.requestJobResult(jobID).get());
		}
	};

	// wait until the job is running
	while (client.listJobs().get().isEmpty()) {
		Thread.sleep(50);
	}

	// kick off resuming
	resumingThread.start();

	// wait for client to connect
	while (resumingThread.getState() != Thread.State.WAITING) {
		Thread.sleep(10);
	}

	// client has connected, we can release the lock
	lock.release();

	resumingThread.sync();
}
 

private void testProgram(
		final MiniClusterWithClientResource cluster,
		final int dataVolumeGb,
		final boolean useForwarder,
		final boolean isSlowSender,
		final boolean isSlowReceiver,
		final int parallelism) throws Exception {
	ClusterClient<?> client = cluster.getClusterClient();

	JobExecutionResult jer = ClientUtils.submitJobAndWaitForResult(
		client,
		createJobGraph(
			dataVolumeGb,
			useForwarder,
			isSlowSender,
			isSlowReceiver,
			parallelism),
		getClass().getClassLoader());

	long dataVolumeMbit = dataVolumeGb * 8192;
	long runtimeSecs = jer.getNetRuntime(TimeUnit.SECONDS);

	int mbitPerSecond = (int) (((double) dataVolumeMbit) / runtimeSecs);

	LOG.info(String.format("Test finished with throughput of %d MBit/s (runtime [secs]: %d, " +
		"data volume [gb/mbits]: %d/%d)", mbitPerSecond, runtimeSecs, dataVolumeGb, dataVolumeMbit));
}
 

/**
 * Tests that a simple job (Source -> Map) with multi regions could restore with operator state.
 *
 * <p>The last subtask of Map function in the 1st stream graph would fail {@code NUM_OF_RESTARTS} times,
 * and it will verify whether the restored state is identical to last completed checkpoint's.
 */
@Test(timeout = 60000)
public void testMultiRegionFailover() {
	try {
		JobGraph jobGraph = createJobGraph();
		ClusterClient<?> client = cluster.getClusterClient();
		ClientUtils.submitJobAndWaitForResult(client, jobGraph, RegionFailoverITCase.class.getClassLoader());
		verifyAfterJobExecuted();
	} catch (Exception e) {
		e.printStackTrace();
		Assert.fail(e.getMessage());
	}
}
 

/**
 * Runs the following program the test program defined in {@link #testProgram(StreamExecutionEnvironment)}
 * followed by the checks in {@link #postSubmit}.
 */
@Test
public void runCheckpointedProgram() throws Exception {
	try {
		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(PARALLELISM);
		env.enableCheckpointing(500);
					env.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 0L));

		testProgram(env);

		JobGraph jobGraph = env.getStreamGraph().getJobGraph();
		try {
			ClientUtils.submitJobAndWaitForResult(cluster.getClusterClient(), jobGraph, getClass().getClassLoader()).getJobExecutionResult();
		} catch (ProgramInvocationException root) {
			Throwable cause = root.getCause();

			// search for nested SuccessExceptions
			int depth = 0;
			while (!(cause instanceof SuccessException)) {
				if (cause == null || depth++ == 20) {
					root.printStackTrace();
					fail("Test failed: " + root.getMessage());
				}
				else {
					cause = cause.getCause();
				}
			}
		}

		postSubmit();
	}
	catch (Exception e) {
		e.printStackTrace();
		Assert.fail(e.getMessage());
	}
}
 

private static String runJobAndGetExternalizedCheckpoint(StateBackend backend, File checkpointDir, @Nullable String externalCheckpoint, ClusterClient<?> client) throws Exception {
	JobGraph initialJobGraph = getJobGraph(backend, externalCheckpoint);
	NotifyingInfiniteTupleSource.countDownLatch = new CountDownLatch(PARALLELISM);

	ClientUtils.submitJob(client, initialJobGraph);

	// wait until all sources have been started
	NotifyingInfiniteTupleSource.countDownLatch.await();

	waitUntilExternalizedCheckpointCreated(checkpointDir, initialJobGraph.getJobID());
	client.cancel(initialJobGraph.getJobID()).get();
	waitUntilCanceled(initialJobGraph.getJobID(), client);

	return getExternalizedCheckpointCheckpointPath(checkpointDir, initialJobGraph.getJobID());
}
 

@Test
public void testJobSubmitCancel() throws Exception {
	TestJobSubmitHandler submitHandler = new TestJobSubmitHandler();
	TestJobCancellationHandler terminationHandler = new TestJobCancellationHandler();
	TestJobExecutionResultHandler testJobExecutionResultHandler =
		new TestJobExecutionResultHandler(
			JobExecutionResultResponseBody.created(new JobResult.Builder()
				.applicationStatus(ApplicationStatus.SUCCEEDED)
				.jobId(jobId)
				.netRuntime(Long.MAX_VALUE)
				.build()));

	try (TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(
		submitHandler,
		terminationHandler,
		testJobExecutionResultHandler)) {

		try (RestClusterClient<?> restClusterClient = createRestClusterClient(restServerEndpoint.getServerAddress().getPort())) {
			Assert.assertFalse(submitHandler.jobSubmitted);
			ClientUtils.submitJobAndWaitForResult(restClusterClient, jobGraph, ClassLoader.getSystemClassLoader());
			Assert.assertTrue(submitHandler.jobSubmitted);

			Assert.assertFalse(terminationHandler.jobCanceled);
			restClusterClient.cancel(jobId).get();
			Assert.assertTrue(terminationHandler.jobCanceled);
		}
	}
}
 

@Test
public void testTriggerSavepointWithCheckpointingDisabled() throws Exception {
	// Config
	final int numTaskManagers = 1;
	final int numSlotsPerTaskManager = 1;

	final Configuration config = new Configuration();

	final MiniClusterWithClientResource cluster = new MiniClusterWithClientResource(
		new MiniClusterResourceConfiguration.Builder()
			.setConfiguration(config)
			.setNumberTaskManagers(numTaskManagers)
			.setNumberSlotsPerTaskManager(numSlotsPerTaskManager)
			.build());
	cluster.before();
	final ClusterClient<?> client = cluster.getClusterClient();

	final JobVertex vertex = new JobVertex("Blocking vertex");
	vertex.setInvokableClass(BlockingNoOpInvokable.class);
	vertex.setParallelism(1);

	final JobGraph graph = new JobGraph(vertex);

	try {
		ClientUtils.submitJob(client, graph);

		client.triggerSavepoint(graph.getJobID(), null).get();

		fail();
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.findThrowable(e, IllegalStateException.class).isPresent());
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, graph.getJobID().toString()).isPresent());
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, "is not a streaming job").isPresent());
	} finally {
		cluster.after();
	}
}
 

protected void runAndCancelJob(Plan plan, final int msecsTillCanceling, int maxTimeTillCanceled) throws Exception {
	// submit job
	final JobGraph jobGraph = getJobGraph(plan);

	final long rpcTimeout = AkkaUtils.getTimeoutAsTime(configuration).toMilliseconds();

	ClusterClient<?> client = CLUSTER.getClusterClient();
	JobSubmissionResult jobSubmissionResult = ClientUtils.submitJob(client, jobGraph);

	Deadline submissionDeadLine = new FiniteDuration(2, TimeUnit.MINUTES).fromNow();

	JobStatus jobStatus = client.getJobStatus(jobSubmissionResult.getJobID()).get(rpcTimeout, TimeUnit.MILLISECONDS);
	while (jobStatus != JobStatus.RUNNING && submissionDeadLine.hasTimeLeft()) {
		Thread.sleep(50);
		jobStatus = client.getJobStatus(jobSubmissionResult.getJobID()).get(rpcTimeout, TimeUnit.MILLISECONDS);
	}
	if (jobStatus != JobStatus.RUNNING) {
		Assert.fail("Job not in state RUNNING.");
	}

	Thread.sleep(msecsTillCanceling);

	client.cancel(jobSubmissionResult.getJobID()).get();

	Deadline cancelDeadline = new FiniteDuration(maxTimeTillCanceled, TimeUnit.MILLISECONDS).fromNow();

	JobStatus jobStatusAfterCancel = client.getJobStatus(jobSubmissionResult.getJobID()).get(rpcTimeout, TimeUnit.MILLISECONDS);
	while (jobStatusAfterCancel != JobStatus.CANCELED && cancelDeadline.hasTimeLeft()) {
		Thread.sleep(50);
		jobStatusAfterCancel = client.getJobStatus(jobSubmissionResult.getJobID()).get(rpcTimeout, TimeUnit.MILLISECONDS);
	}
	if (jobStatusAfterCancel != JobStatus.CANCELED) {
		Assert.fail("Failed to cancel job with ID " + jobSubmissionResult.getJobID() + '.');
	}
}
 

@Test
public void testJobSubmissionFailureThrowsProgramInvocationException() throws Exception {
	try (final TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(new SubmissionFailingHandler())) {
		RestClusterClient<?> restClusterClient = createRestClusterClient(restServerEndpoint.getServerAddress().getPort());

		try {
			ClientUtils.submitJobAndWaitForResult(restClusterClient, jobGraph, ClassLoader.getSystemClassLoader());
		} catch (final ProgramInvocationException expected) {
			// expected
		} finally {
			restClusterClient.close();
		}
	}
}
 

/**
 * Verify operators would be notified as checkpoint aborted.
 *
 * <p>The job would run with at least two checkpoints. The 1st checkpoint would fail due to add checkpoint to store,
 * and the 2nd checkpoint would decline by async checkpoint phase of 'DeclineSink'.
 *
 * <p>The job graph looks like:
 * NormalSource --> keyBy --> NormalMap --> DeclineSink
 */
@Test(timeout = TEST_TIMEOUT)
public void testNotifyCheckpointAborted() throws Exception {
	final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.enableCheckpointing(200, CheckpointingMode.EXACTLY_ONCE);
	env.getCheckpointConfig().enableUnalignedCheckpoints(unalignedCheckpointEnabled);
	env.getCheckpointConfig().setTolerableCheckpointFailureNumber(1);
	env.disableOperatorChaining();
	env.setParallelism(1);

	final StateBackend failingStateBackend = new DeclineSinkFailingStateBackend(checkpointPath);
	env.setStateBackend(failingStateBackend);

	env.addSource(new NormalSource()).name("NormalSource")
		.keyBy((KeySelector<Tuple2<Integer, Integer>, Integer>) value -> value.f0)
		.transform("NormalMap", TypeInformation.of(Integer.class), new NormalMap())
		.transform(DECLINE_SINK_NAME, TypeInformation.of(Object.class), new DeclineSink());

	final ClusterClient<?> clusterClient = cluster.getClusterClient();
	JobGraph jobGraph = env.getStreamGraph().getJobGraph();
	JobID jobID = jobGraph.getJobID();

	ClientUtils.submitJob(clusterClient, jobGraph);

	TestingCompletedCheckpointStore.addCheckpointLatch.await();
	TestingCompletedCheckpointStore.abortCheckpointLatch.trigger();

	verifyAllOperatorsNotifyAborted();
	resetAllOperatorsNotifyAbortedLatches();
	verifyAllOperatorsNotifyAbortedTimes(1);

	DeclineSink.waitLatch.trigger();
	verifyAllOperatorsNotifyAborted();
	verifyAllOperatorsNotifyAbortedTimes(2);

	clusterClient.cancel(jobID).get();
}
 

/**
 * Tests simple value state queryable state instance with a default value
 * set. Each source emits (subtaskIndex, 0)..(subtaskIndex, numElements)
 * tuples, the key is mapped to 1 but key 0 is queried which should throw
 * a {@link UnknownKeyOrNamespaceException} exception.
 *
 * @throws UnknownKeyOrNamespaceException thrown due querying a non-existent key
 */
@Test(expected = UnknownKeyOrNamespaceException.class)
public void testValueStateDefault() throws Throwable {
	final Deadline deadline = Deadline.now().plus(TEST_TIMEOUT);
	final long numElements = 1024L;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setStateBackend(stateBackend);
	env.setParallelism(maxParallelism);
	// Very important, because cluster is shared between tests and we
	// don't explicitly check that all slots are available before
	// submitting.
	env.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000L));

	DataStream<Tuple2<Integer, Long>> source = env.addSource(new TestAscendingValueSource(numElements));

	ValueStateDescriptor<Tuple2<Integer, Long>> valueState = new ValueStateDescriptor<>(
			"any", source.getType(), 	Tuple2.of(0, 1337L));

	// only expose key "1"
	QueryableStateStream<Integer, Tuple2<Integer, Long>> queryableState = source.keyBy(
			new KeySelector<Tuple2<Integer, Long>, Integer>() {
				private static final long serialVersionUID = 4509274556892655887L;

				@Override
				public Integer getKey(Tuple2<Integer, Long> value) {
					return 1;
				}
			}).asQueryableState("hakuna", valueState);

	try (AutoCancellableJob autoCancellableJob = new AutoCancellableJob(deadline, clusterClient, env)) {

		final JobID jobId = autoCancellableJob.getJobId();
		final JobGraph jobGraph = autoCancellableJob.getJobGraph();

		ClientUtils.submitJob(clusterClient, jobGraph);

		// Now query
		int key = 0;
		CompletableFuture<ValueState<Tuple2<Integer, Long>>> future = getKvState(
				deadline,
				client,
				jobId,
				queryableState.getQueryableStateName(),
				key,
				BasicTypeInfo.INT_TYPE_INFO,
				valueState,
				true,
				executor);

		try {
			future.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
		} catch (ExecutionException | CompletionException e) {
			// get() on a completedExceptionally future wraps the
			// exception in an ExecutionException.
			throw e.getCause();
		}
	}
}
 

@Test
public void testSubmitJobAndWaitForExecutionResult() throws Exception {
	final TestJobExecutionResultHandler testJobExecutionResultHandler =
		new TestJobExecutionResultHandler(
			new RestHandlerException("should trigger retry", HttpResponseStatus.SERVICE_UNAVAILABLE),
			JobExecutionResultResponseBody.inProgress(),
			JobExecutionResultResponseBody.created(new JobResult.Builder()
				.applicationStatus(ApplicationStatus.SUCCEEDED)
				.jobId(jobId)
				.netRuntime(Long.MAX_VALUE)
				.accumulatorResults(Collections.singletonMap("testName", new SerializedValue<>(OptionalFailure.of(1.0))))
				.build()),
			JobExecutionResultResponseBody.created(new JobResult.Builder()
				.applicationStatus(ApplicationStatus.FAILED)
				.jobId(jobId)
				.netRuntime(Long.MAX_VALUE)
				.serializedThrowable(new SerializedThrowable(new RuntimeException("expected")))
				.build()));

	// fail first HTTP polling attempt, which should not be a problem because of the retries
	final AtomicBoolean firstPollFailed = new AtomicBoolean();
	failHttpRequest = (messageHeaders, messageParameters, requestBody) ->
		messageHeaders instanceof JobExecutionResultHeaders && !firstPollFailed.getAndSet(true);

	try (TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(
		testJobExecutionResultHandler,
		new TestJobSubmitHandler())) {
		RestClusterClient<?> restClusterClient = createRestClusterClient(restServerEndpoint.getServerAddress().getPort());

		try {
			JobExecutionResult jobExecutionResult;

			jobExecutionResult = ClientUtils.submitJobAndWaitForResult(restClusterClient, jobGraph, ClassLoader.getSystemClassLoader());
			assertThat(jobExecutionResult.getJobID(), equalTo(jobId));
			assertThat(jobExecutionResult.getNetRuntime(), equalTo(Long.MAX_VALUE));
			assertThat(
				jobExecutionResult.getAllAccumulatorResults(),
				equalTo(Collections.singletonMap("testName", 1.0)));

			try {
				ClientUtils.submitJobAndWaitForResult(restClusterClient, jobGraph, ClassLoader.getSystemClassLoader());
				fail("Expected exception not thrown.");
			} catch (final ProgramInvocationException e) {
				final Optional<RuntimeException> cause = ExceptionUtils.findThrowable(e, RuntimeException.class);

				assertThat(cause.isPresent(), is(true));
				assertThat(cause.get().getMessage(), equalTo("expected"));
			}
		} finally {
			restClusterClient.close();
		}
	}
}
 

/**
 * Similar tests as {@link #testValueState()} but before submitting the
 * job, we already issue one request which fails.
 */
@Test
public void testQueryNonStartedJobState() throws Exception {
	final Deadline deadline = Deadline.now().plus(TEST_TIMEOUT);
	final long numElements = 1024L;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setStateBackend(stateBackend);
	env.setParallelism(maxParallelism);
	// Very important, because clusterClient is shared between tests and we
	// don't explicitly check that all slots are available before
	// submitting.
	env.setRestartStrategy(RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, 1000L));

	DataStream<Tuple2<Integer, Long>> source = env.addSource(new TestAscendingValueSource(numElements));

	ValueStateDescriptor<Tuple2<Integer, Long>> valueState = new ValueStateDescriptor<>(
		"any", source.getType(), 	null);

	QueryableStateStream<Integer, Tuple2<Integer, Long>> queryableState =
			source.keyBy(new KeySelector<Tuple2<Integer, Long>, Integer>() {

				private static final long serialVersionUID = 7480503339992214681L;

				@Override
				public Integer getKey(Tuple2<Integer, Long> value) {
					return value.f0;
				}
			}).asQueryableState("hakuna", valueState);

	try (AutoCancellableJob autoCancellableJob = new AutoCancellableJob(deadline, clusterClient, env)) {

		final JobID jobId = autoCancellableJob.getJobId();
		final JobGraph jobGraph = autoCancellableJob.getJobGraph();

		long expected = numElements;

		// query once
		client.getKvState(
				autoCancellableJob.getJobId(),
				queryableState.getQueryableStateName(),
				0,
				BasicTypeInfo.INT_TYPE_INFO,
				valueState);

		ClientUtils.submitJob(clusterClient, jobGraph);
		executeValueQuery(deadline, client, jobId, "hakuna", valueState, expected);
	}
}
 

/**
 * Tests that metrics registered on the JobManager are actually accessible via JMX.
 */
@Test
public void testJobManagerJMXMetricAccess() throws Exception {
	Deadline deadline = Deadline.now().plus(Duration.ofMinutes(2));

	try {
		JobVertex sourceJobVertex = new JobVertex("Source");
		sourceJobVertex.setInvokableClass(BlockingInvokable.class);

		JobGraph jobGraph = new JobGraph("TestingJob", sourceJobVertex);
		jobGraph.setSnapshotSettings(new JobCheckpointingSettings(
			Collections.<JobVertexID>emptyList(),
			Collections.<JobVertexID>emptyList(),
			Collections.<JobVertexID>emptyList(),
			new CheckpointCoordinatorConfiguration(
				500,
				500,
				50,
				5,
				CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
				true,
				false,
				false,
				0),
			null));

		ClusterClient<?> client = MINI_CLUSTER_RESOURCE.getClusterClient();
		ClientUtils.submitJob(client, jobGraph);

		FutureUtils.retrySuccessfulWithDelay(
			() -> client.getJobStatus(jobGraph.getJobID()),
			Time.milliseconds(10),
			deadline,
			status -> status == JobStatus.RUNNING,
			TestingUtils.defaultScheduledExecutor()
		).get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);

		MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
		Set<ObjectName> nameSet = mBeanServer.queryNames(new ObjectName("org.apache.flink.jobmanager.job.lastCheckpointSize:job_name=TestingJob,*"), null);
		Assert.assertEquals(1, nameSet.size());
		assertEquals(-1L, mBeanServer.getAttribute(nameSet.iterator().next(), "Value"));

		BlockingInvokable.unblock();
	} finally {
		BlockingInvokable.unblock();
	}
}