org.apache.flink.runtime.messages.Acknowledge Java Examples

@Override
public CompletableFuture<Acknowledge> notifyKvStateRegistered(
		final JobID jobId,
		final JobVertexID jobVertexId,
		final KeyGroupRange keyGroupRange,
		final String registrationName,
		final KvStateID kvStateId,
		final InetSocketAddress kvStateServerAddress) {

	try {
		schedulerNG.notifyKvStateRegistered(jobId, jobVertexId, keyGroupRange, registrationName, kvStateId, kvStateServerAddress);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} catch (FlinkJobNotFoundException e) {
		log.info("Error while receiving notification about key-value state registration", e);
		return FutureUtils.completedExceptionally(e);
	}
}
 

@Test
public void testDisposeSavepointSuccess() throws Exception {
	replaceStdOutAndStdErr();

	String savepointPath = "expectedSavepointPath";

	ClusterClient clusterClient = new DisposeSavepointClusterClient(
		(String path) -> CompletableFuture.completedFuture(Acknowledge.get()), getConfiguration());

	try {

		CliFrontend frontend = new MockedCliFrontend(clusterClient);

		String[] parameters = { "-d", savepointPath };
		frontend.savepoint(parameters);

		String outMsg = buffer.toString();
		assertTrue(outMsg.contains(savepointPath));
		assertTrue(outMsg.contains("disposed"));
	}
	finally {
		clusterClient.shutdown();
		restoreStdOutAndStdErr();
	}
}
 

@Test
public void testSuccessfulJobSubmission() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Collections.singleton(jobGraphFile.toFile())), mockGateway)
		.get();
}
 

/**
 * Sends a {@link JobManagerMessages.DisposeSavepoint} message to the job manager.
 */
private void disposeSavepoint(ClusterClient<?> clusterClient, String savepointPath) throws FlinkException {
	Preconditions.checkNotNull(savepointPath, "Missing required argument: savepoint path. " +
		"Usage: bin/flink savepoint -d <savepoint-path>");

	logAndSysout("Disposing savepoint '" + savepointPath + "'.");

	final CompletableFuture<Acknowledge> disposeFuture = clusterClient.disposeSavepoint(savepointPath);

	logAndSysout("Waiting for response...");

	try {
		disposeFuture.get(clientTimeout.toMillis(), TimeUnit.MILLISECONDS);
	} catch (Exception e) {
		throw new FlinkException("Disposing the savepoint '" + savepointPath + "' failed.", e);
	}

	logAndSysout("Savepoint '" + savepointPath + "' disposed.");
}
 

/**
 * Tests that the stopping RPC call is sent upon stopping requests.
 */
@Test
public void testStopRpc() throws Exception {
	final JobID jid = new JobID();
	final JobVertex vertex = new JobVertex("vertex");
	vertex.setInvokableClass(NoOpInvokable.class);
	vertex.setParallelism(5);

	final ExecutionGraph graph = ExecutionGraphTestUtils.createSimpleTestGraph(jid, vertex);
	final Execution exec = graph.getJobVertex(vertex.getID()).getTaskVertices()[0].getCurrentExecutionAttempt();

	final TaskManagerGateway gateway = mock(TaskManagerGateway.class);
	when(gateway.submitTask(any(TaskDeploymentDescriptor.class), any(Time.class)))
			.thenReturn(CompletableFuture.completedFuture(Acknowledge.get()));
	when(gateway.stopTask(any(ExecutionAttemptID.class), any(Time.class)))
			.thenReturn(CompletableFuture.completedFuture(Acknowledge.get()));

	final SimpleSlot slot = ExecutionGraphTestUtils.createMockSimpleSlot(gateway);

	exec.tryAssignResource(slot);
	exec.deploy();
	exec.switchToRunning();
	assertEquals(ExecutionState.RUNNING, exec.getState());

	exec.stop();
	assertEquals(ExecutionState.RUNNING, exec.getState());

	verify(gateway, times(1)).stopTask(any(ExecutionAttemptID.class), any(Time.class));

	exec.markFinished();
	assertEquals(ExecutionState.FINISHED, exec.getState());
}
 

public TestingRestfulGateway(
		String address,
		String hostname,
		Function<JobID, CompletableFuture<Acknowledge>> cancelJobFunction,
		Function<JobID, CompletableFuture<Acknowledge>> stopJobFunction,
		Function<JobID, CompletableFuture<ArchivedExecutionGraph>> requestJobFunction,
		Function<JobID, CompletableFuture<JobResult>> requestJobResultFunction,
		Function<JobID, CompletableFuture<JobStatus>> requestJobStatusFunction,
		Supplier<CompletableFuture<MultipleJobsDetails>> requestMultipleJobDetailsSupplier,
		Supplier<CompletableFuture<ClusterOverview>> requestClusterOverviewSupplier,
		Supplier<CompletableFuture<Collection<String>>> requestMetricQueryServicePathsSupplier,
		Supplier<CompletableFuture<Collection<Tuple2<ResourceID, String>>>> requestTaskManagerMetricQueryServicePathsSupplier,
		BiFunction<JobID, JobVertexID, CompletableFuture<OperatorBackPressureStatsResponse>> requestOperatorBackPressureStatsFunction,
		BiFunction<JobID, String, CompletableFuture<String>> triggerSavepointFunction) {
	this.address = address;
	this.hostname = hostname;
	this.cancelJobFunction = cancelJobFunction;
	this.stopJobFunction = stopJobFunction;
	this.requestJobFunction = requestJobFunction;
	this.requestJobResultFunction = requestJobResultFunction;
	this.requestJobStatusFunction = requestJobStatusFunction;
	this.requestMultipleJobDetailsSupplier = requestMultipleJobDetailsSupplier;
	this.requestClusterOverviewSupplier = requestClusterOverviewSupplier;
	this.requestMetricQueryServicePathsSupplier = requestMetricQueryServicePathsSupplier;
	this.requestTaskManagerMetricQueryServicePathsSupplier = requestTaskManagerMetricQueryServicePathsSupplier;
	this.requestOperatorBackPressureStatsFunction = requestOperatorBackPressureStatsFunction;
	this.triggerSavepointFunction = triggerSavepointFunction;
}
 

@Test
public void testRequestKvStateWithoutRegistration() throws Exception {
	final JobGraph graph = createKvJobGraph();

	final JobMaster jobMaster = createJobMaster(
		configuration,
		graph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);
	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

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

		// lookup location
		try {
			jobMasterGateway.requestKvStateLocation(graph.getJobID(), "unknown").get();
			fail("Expected to fail with UnknownKvStateLocation");
		} catch (Exception e) {
			assertTrue(ExceptionUtils.findThrowable(e, UnknownKvStateLocation.class).isPresent());
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Tests that the uploaded blobs are being cleaned up in case of a job submission failure.
 */
@Test
public void testBlobServerCleanupWhenJobSubmissionFails() throws Exception {
	failJobMasterCreationWith.set(() -> new FlinkException("Test exception."));
	final CompletableFuture<Acknowledge> submissionFuture = dispatcherGateway.submitJob(jobGraph, timeout);

	try {
		submissionFuture.get();
		fail("Job submission was expected to fail.");
	} catch (ExecutionException ee) {
		assertThat(ExceptionUtils.findThrowable(ee, JobSubmissionException.class).isPresent(), is(true));
	}

	assertThatHABlobsHaveBeenRemoved();
}
 

public CompletableFuture<Acknowledge> setFencingTokenInMainThread(UUID fencingToken, Time timeout) {
	return callAsyncWithoutFencing(
		() -> {
			setFencingToken(fencingToken);

			return Acknowledge.get();
		},
		timeout);
}
 

@Test
public void testRejectionOnCountMismatch() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}
	final Path countExceedingFile = TEMPORARY_FOLDER.newFile().toPath();

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList(jobGraphFile.toFile(), countExceedingFile.toFile())), mockGateway)
			.get();
	} catch (Exception e) {
		ExceptionUtils.findThrowable(e, candidate -> candidate instanceof RestHandlerException && candidate.getMessage().contains("count"));
	}
}
 

/**
 * Tests that the {@link JobManagerRunner} always waits for the previous leadership operation
 * (granting or revoking leadership) to finish before starting a new leadership operation.
 */
@Test
public void testConcurrentLeadershipOperationsBlockingSuspend() throws Exception {
	final CompletableFuture<Acknowledge> suspendedFuture = new CompletableFuture<>();

	TestingJobMasterServiceFactory jobMasterServiceFactory = new TestingJobMasterServiceFactory(
		() -> new TestingJobMasterService(
			"localhost",
			e -> suspendedFuture));
	JobManagerRunner jobManagerRunner = createJobManagerRunner(jobMasterServiceFactory);

	jobManagerRunner.start();

	leaderElectionService.isLeader(UUID.randomUUID()).get();

	leaderElectionService.notLeader();

	final CompletableFuture<UUID> leaderFuture = leaderElectionService.isLeader(UUID.randomUUID());

	// the new leadership should wait first for the suspension to happen
	assertThat(leaderFuture.isDone(), is(false));

	try {
		leaderFuture.get(1L, TimeUnit.MILLISECONDS);
		fail("Granted leadership even though the JobMaster has not been suspended.");
	} catch (TimeoutException expected) {
		// expected
	}

	suspendedFuture.complete(Acknowledge.get());

	leaderFuture.get();
}
 

@Test
public void testRequestKvStateWithIrrelevantRegistration() throws Exception {
	final JobGraph graph = createKvJobGraph();

	final JobMaster jobMaster = createJobMaster(
		configuration,
		graph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);
	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

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

		// register an irrelevant KvState
		try {
			jobMasterGateway.notifyKvStateRegistered(
				new JobID(),
				new JobVertexID(),
				new KeyGroupRange(0, 0),
				"any-name",
				new KvStateID(),
				new InetSocketAddress(InetAddress.getLocalHost(), 1233)).get();
			fail("Expected to fail with FlinkJobNotFoundException.");
		} catch (Exception e) {
			assertTrue(ExceptionUtils.findThrowable(e, FlinkJobNotFoundException.class).isPresent());
		}
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Cleanup application and shut down cluster.
 *
 * @param finalStatus of the Flink application
 * @param diagnostics diagnostics message for the Flink application or {@code null}
 */
@Override
public CompletableFuture<Acknowledge> deregisterApplication(
		final ApplicationStatus finalStatus,
		@Nullable final String diagnostics) {
	log.info("Shut down cluster because application is in {}, diagnostics {}.", finalStatus, diagnostics);

	try {
		internalDeregisterApplication(finalStatus, diagnostics);
	} catch (ResourceManagerException e) {
		log.warn("Could not properly shutdown the application.", e);
	}

	return CompletableFuture.completedFuture(Acknowledge.get());
}
 

@Override
public CompletableFuture<Acknowledge> freeSlot(AllocationID allocationId, Throwable cause, Time timeout) {
	final BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>> currentFreeSlotFunction = freeSlotFunction;

	if (currentFreeSlotFunction != null) {
		return currentFreeSlotFunction.apply(allocationId, cause);
	} else {
		return CompletableFuture.completedFuture(Acknowledge.get());
	}
}
 

DisposeSavepointClusterClient(
	Function<String, CompletableFuture<Acknowledge>> disposeSavepointFunction,
	Configuration configuration) throws Exception {
	super(configuration, StandaloneClusterId.getInstance());

	this.disposeSavepointFunction = Preconditions.checkNotNull(disposeSavepointFunction);
}
 

@Override
public CompletableFuture<Acknowledge> submitJob(JobGraph jobGraph, Time timeout) {
	log.info("Received JobGraph submission {} ({}).", jobGraph.getJobID(), jobGraph.getName());

	try {
		if (isDuplicateJob(jobGraph.getJobID())) {
			return FutureUtils.completedExceptionally(
				new JobSubmissionException(jobGraph.getJobID(), "Job has already been submitted."));
		} else {
			return internalSubmitJob(jobGraph);
		}
	} catch (FlinkException e) {
		return FutureUtils.completedExceptionally(e);
	}
}
 

/**
 * Rejects the pending slot request by failing the request future with a
 * {@link SlotAllocationException}.
 *
 * @param pendingSlotRequest to reject
 * @param cause of the rejection
 */
private void rejectPendingSlotRequest(PendingSlotRequest pendingSlotRequest, Exception cause) {
	CompletableFuture<Acknowledge> request = pendingSlotRequest.getRequestFuture();

	if (null != request) {
		request.completeExceptionally(new SlotAllocationException(cause));
	} else {
		LOG.debug("Cannot reject pending slot request {}, since no request has been sent.", pendingSlotRequest.getAllocationId());
	}
}
 

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

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

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

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

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

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

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

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

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

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

@Override
public CompletableFuture<Acknowledge> rescaleJob(JobID jobId, int newParallelism) {

	final RescalingTriggerHeaders rescalingTriggerHeaders = RescalingTriggerHeaders.getInstance();
	final RescalingTriggerMessageParameters rescalingTriggerMessageParameters = rescalingTriggerHeaders.getUnresolvedMessageParameters();
	rescalingTriggerMessageParameters.jobPathParameter.resolve(jobId);
	rescalingTriggerMessageParameters.rescalingParallelismQueryParameter.resolve(Collections.singletonList(newParallelism));

	final CompletableFuture<TriggerResponse> rescalingTriggerResponseFuture = sendRequest(
		rescalingTriggerHeaders,
		rescalingTriggerMessageParameters);

	final CompletableFuture<AsynchronousOperationInfo> rescalingOperationFuture = rescalingTriggerResponseFuture.thenCompose(
		(TriggerResponse triggerResponse) -> {
			final TriggerId triggerId = triggerResponse.getTriggerId();
			final RescalingStatusHeaders rescalingStatusHeaders = RescalingStatusHeaders.getInstance();
			final RescalingStatusMessageParameters rescalingStatusMessageParameters = rescalingStatusHeaders.getUnresolvedMessageParameters();

			rescalingStatusMessageParameters.jobPathParameter.resolve(jobId);
			rescalingStatusMessageParameters.triggerIdPathParameter.resolve(triggerId);

			return pollResourceAsync(
				() -> sendRequest(
					rescalingStatusHeaders,
					rescalingStatusMessageParameters));
		});

	return rescalingOperationFuture.thenApply(
		(AsynchronousOperationInfo asynchronousOperationInfo) -> {
			if (asynchronousOperationInfo.getFailureCause() == null) {
				return Acknowledge.get();
			} else {
				throw new CompletionException(asynchronousOperationInfo.getFailureCause());
			}
		});
}
 

@Override
public CompletableFuture<Acknowledge> setNewFencingToken(UUID fencingToken, Time timeout) {
	enteringSetNewFencingToken.trigger();
	try {
		triggerSetNewFencingToken.await();
	} catch (InterruptedException e) {
		throw new RuntimeException("TriggerSetNewFencingToken OneShotLatch was interrupted.");
	}

	setFencingToken(fencingToken);

	return CompletableFuture.completedFuture(Acknowledge.get());
}
 

/**
 * Updates the task execution state for a given task.
 *
 * @param taskExecutionState New task execution state for a given task
 * @return Acknowledge the task execution state update
 */
@Override
public CompletableFuture<Acknowledge> updateTaskExecutionState(
		final TaskExecutionState taskExecutionState) {
	checkNotNull(taskExecutionState, "taskExecutionState");

	if (executionGraph.updateState(taskExecutionState)) {
		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		return FutureUtils.completedExceptionally(
			new ExecutionGraphException("The execution attempt " +
				taskExecutionState.getID() + " was not found."));
	}
}
 

@Override
public CompletableFuture<Acknowledge> requestSlot(
		JobMasterId jobMasterId,
		SlotRequest slotRequest,
		final Time timeout) {

	JobID jobId = slotRequest.getJobId();
	JobManagerRegistration jobManagerRegistration = jobManagerRegistrations.get(jobId);

	if (null != jobManagerRegistration) {
		if (Objects.equals(jobMasterId, jobManagerRegistration.getJobMasterId())) {
			log.info("Request slot with profile {} for job {} with allocation id {}.",
				slotRequest.getResourceProfile(),
				slotRequest.getJobId(),
				slotRequest.getAllocationId());

			try {
				slotManager.registerSlotRequest(slotRequest);
			} catch (ResourceManagerException e) {
				return FutureUtils.completedExceptionally(e);
			}

			return CompletableFuture.completedFuture(Acknowledge.get());
		} else {
			return FutureUtils.completedExceptionally(new ResourceManagerException("The job leader's id " +
				jobManagerRegistration.getJobMasterId() + " does not match the received id " + jobMasterId + '.'));
		}

	} else {
		return FutureUtils.completedExceptionally(new ResourceManagerException("Could not find registered job manager for job " + jobId + '.'));
	}
}
 

public TestingDispatcherGateway(
		String address,
		String hostname,
		Function<JobID, CompletableFuture<Acknowledge>> cancelJobFunction,
		Function<JobID, CompletableFuture<Acknowledge>> stopJobFunction,
		Function<JobID, CompletableFuture<ArchivedExecutionGraph>> requestJobFunction,
		Function<JobID, CompletableFuture<JobResult>> requestJobResultFunction,
		Function<JobID, CompletableFuture<JobStatus>> requestJobStatusFunction,
		Supplier<CompletableFuture<MultipleJobsDetails>> requestMultipleJobDetailsSupplier,
		Supplier<CompletableFuture<ClusterOverview>> requestClusterOverviewSupplier,
		Supplier<CompletableFuture<Collection<String>>> requestMetricQueryServicePathsSupplier,
		Supplier<CompletableFuture<Collection<Tuple2<ResourceID, String>>>> requestTaskManagerMetricQueryServicePathsSupplier,
		BiFunction<JobID, JobVertexID, CompletableFuture<OperatorBackPressureStatsResponse>> requestOperatorBackPressureStatsFunction,
		BiFunction<JobID, String, CompletableFuture<String>> triggerSavepointFunction,
		Function<JobGraph, CompletableFuture<Acknowledge>> submitFunction,
		Supplier<CompletableFuture<Collection<JobID>>> listFunction,
		int blobServerPort,
		DispatcherId fencingToken,
		Function<JobID, CompletableFuture<ArchivedExecutionGraph>> requestArchivedJobFunction) {
	super(
		address,
		hostname,
		cancelJobFunction,
		stopJobFunction,
		requestJobFunction,
		requestJobResultFunction,
		requestJobStatusFunction,
		requestMultipleJobDetailsSupplier,
		requestClusterOverviewSupplier,
		requestMetricQueryServicePathsSupplier,
		requestTaskManagerMetricQueryServicePathsSupplier,
		requestOperatorBackPressureStatsFunction,
		triggerSavepointFunction);
	this.submitFunction = submitFunction;
	this.listFunction = listFunction;
	this.blobServerPort = blobServerPort;
	this.fencingToken = fencingToken;
	this.requestArchivedJobFunction = requestArchivedJobFunction;
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			Arrays.asList(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN),
			timerService,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

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

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

	try {
		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report
 * fails.
 */
@Test
public void testInitialSlotReportFailure() throws Exception {
	final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
		.setTaskSlotTable(taskSlotTable)
		.setTaskManagerLocation(taskManagerLocation)
		.build();
	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	taskExecutor.start();

	try {
		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
		testingResourceManagerGateway.setSendSlotReportFunction(
			resourceIDInstanceIDSlotReportTuple3 -> {
				try {
					return responseQueue.take();
				} catch (InterruptedException e) {
					return FutureUtils.completedExceptionally(e);
				}
			});

		final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				testingResourceManagerGateway.getOwnResourceId(),
				new ClusterInformation("foobar", 1234)));

		final CountDownLatch numberRegistrations = new CountDownLatch(2);

		testingResourceManagerGateway.setRegisterTaskExecutorFunction(new Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>>() {
			@Override
			public CompletableFuture<RegistrationResponse> apply(Tuple4<String, ResourceID, Integer, HardwareDescription> stringResourceIDIntegerHardwareDescriptionTuple4) {
				numberRegistrations.countDown();
				return registrationResponse;
			}
		});

		responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
		responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));

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

		//wait for the second registration attempt
		numberRegistrations.await();
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

@Override
public CompletableFuture<Acknowledge> submitTask(TaskDeploymentDescriptor tdd, Time timeout) {
	CompletableFuture<Acknowledge> future = new CompletableFuture<>();
	future.completeExceptionally(new Exception(ERROR_MESSAGE));
	return future;
}
 

@Override
public CompletableFuture<Acknowledge> triggerCheckpoint(ExecutionAttemptID executionAttemptID, long checkpointID, long checkpointTimestamp, CheckpointOptions checkpointOptions) {
	return CompletableFuture.completedFuture(Acknowledge.get());
}
 

@Override
public CompletableFuture<Acknowledge> updateTaskExecutionState(TaskExecutionState taskExecutionState) {
	return updateTaskExecutionStateFunction.apply(taskExecutionState);
}
 

private void runJobFailureWhenTaskExecutorTerminatesTest(
		HeartbeatServices heartbeatServices,
		BiConsumer<LocalTaskManagerLocation, JobMasterGateway> jobReachedRunningState,
		BiFunction<JobMasterGateway, ResourceID, BiConsumer<ResourceID, AllocatedSlotReport>> heartbeatConsumerFunction) throws Exception {
	final JobGraph jobGraph = createSingleVertexJobGraph();
	final TestingOnCompletionActions onCompletionActions = new TestingOnCompletionActions();
	final JobMaster jobMaster = createJobMaster(
		new Configuration(),
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build(),
		heartbeatServices,
		onCompletionActions);

	try {
		jobMaster.start(jobMasterId).get();

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

		final LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
		final CompletableFuture<ExecutionAttemptID> taskDeploymentFuture = new CompletableFuture<>();
		final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setSubmitTaskConsumer((taskDeploymentDescriptor, jobMasterId) -> {
				taskDeploymentFuture.complete(taskDeploymentDescriptor.getExecutionAttemptId());
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setHeartbeatJobManagerConsumer(heartbeatConsumerFunction.apply(jobMasterGateway, taskManagerLocation.getResourceID()))
			.createTestingTaskExecutorGateway();

		final Collection<SlotOffer> slotOffers = registerSlotsAtJobMaster(1, jobMasterGateway, taskExecutorGateway, taskManagerLocation);
		assertThat(slotOffers, hasSize(1));

		final ExecutionAttemptID executionAttemptId = taskDeploymentFuture.get();

		jobMasterGateway.updateTaskExecutionState(new TaskExecutionState(jobGraph.getJobID(), executionAttemptId, ExecutionState.RUNNING)).get();

		jobReachedRunningState.accept(taskManagerLocation, jobMasterGateway);

		final ArchivedExecutionGraph archivedExecutionGraph = onCompletionActions.getJobReachedGloballyTerminalStateFuture().get();

		assertThat(archivedExecutionGraph.getState(), is(JobStatus.FAILED));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

public Builder setSubmitFunction(Function<JobGraph, CompletableFuture<Acknowledge>> submitFunction) {
	this.submitFunction = submitFunction;
	return this;
}