Java Code Examples for org.apache.flink.runtime.concurrent.FutureUtils#completedExceptionally()

/**
 * Tests that the send operation is being retried.
 */
@Test
public void testRetriableSendOperationIfConnectionErrorOrServiceUnavailable() throws Exception {
	final PingRestHandler pingRestHandler = new PingRestHandler(
		FutureUtils.completedExceptionally(new RestHandlerException("test exception", HttpResponseStatus.SERVICE_UNAVAILABLE)),
		CompletableFuture.completedFuture(EmptyResponseBody.getInstance()));

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

		try {
			final AtomicBoolean firstPollFailed = new AtomicBoolean();
			failHttpRequest = (messageHeaders, messageParameters, requestBody) ->
				messageHeaders instanceof PingRestHandlerHeaders && !firstPollFailed.getAndSet(true);

			restClusterClient.sendRequest(PingRestHandlerHeaders.INSTANCE).get();
		} finally {
			restClusterClient.shutdown();
		}
	}
}
 

@Override
public CompletableFuture<Object> updateGlobalAggregate(String aggregateName, Object aggregand, byte[] serializedAggregateFunction) {

	AggregateFunction aggregateFunction = null;
	try {
		aggregateFunction = InstantiationUtil.deserializeObject(serializedAggregateFunction, userCodeLoader);
	} catch (Exception e) {
		log.error("Error while attempting to deserialize user AggregateFunction.");
		return FutureUtils.completedExceptionally(e);
	}

	Object accumulator = accumulators.get(aggregateName);
	if(null == accumulator) {
		accumulator = aggregateFunction.createAccumulator();
	}
	accumulator = aggregateFunction.add(aggregand, accumulator);
	accumulators.put(aggregateName, accumulator);
	return CompletableFuture.completedFuture(aggregateFunction.getResult(accumulator));
}
 

@Override
public CompletableFuture<Acknowledge> confirmCheckpoint(
		ExecutionAttemptID executionAttemptID,
		long checkpointId,
		long checkpointTimestamp) {
	log.debug("Confirm checkpoint {}@{} for {}.", checkpointId, checkpointTimestamp, executionAttemptID);

	final Task task = taskSlotTable.getTask(executionAttemptID);

	if (task != null) {
		task.notifyCheckpointComplete(checkpointId);

		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		final String message = "TaskManager received a checkpoint confirmation for unknown task " + executionAttemptID + '.';

		log.debug(message);
		return FutureUtils.completedExceptionally(new CheckpointException(message));
	}
}
 

@VisibleForTesting
public CompletableFuture<CompletedCheckpoint> triggerCheckpoint(
		CheckpointProperties props,
		@Nullable String externalSavepointLocation,
		boolean isPeriodic,
		boolean advanceToEndOfTime) {

	if (advanceToEndOfTime && !(props.isSynchronous() && props.isSavepoint())) {
		return FutureUtils.completedExceptionally(new IllegalArgumentException(
			"Only synchronous savepoints are allowed to advance the watermark to MAX."));
	}

	CheckpointTriggerRequest request = new CheckpointTriggerRequest(props, externalSavepointLocation, isPeriodic, advanceToEndOfTime);
	requestDecider
		.chooseRequestToExecute(request, isTriggering, lastCheckpointCompletionRelativeTime)
		.ifPresent(this::startTriggeringCheckpoint);
	return request.onCompletionPromise;
}
 

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

/**
 * Tests that the send operation is being retried.
 */
@Test
public void testRetriableSendOperationIfConnectionErrorOrServiceUnavailable() throws Exception {
	final PingRestHandler pingRestHandler = new PingRestHandler(
		FutureUtils.completedExceptionally(new RestHandlerException("test exception", HttpResponseStatus.SERVICE_UNAVAILABLE)),
		CompletableFuture.completedFuture(EmptyResponseBody.getInstance()));

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

		try {
			final AtomicBoolean firstPollFailed = new AtomicBoolean();
			failHttpRequest = (messageHeaders, messageParameters, requestBody) ->
				messageHeaders instanceof PingRestHandlerHeaders && !firstPollFailed.getAndSet(true);

			restClusterClient.sendRequest(PingRestHandlerHeaders.INSTANCE).get();
		} finally {
			restClusterClient.close();
		}
	}
}
 

@Override
public CompletableFuture<Object> updateGlobalAggregate(String aggregateName, Object aggregand, byte[] serializedAggregateFunction) {

	AggregateFunction aggregateFunction = null;
	try {
		aggregateFunction = InstantiationUtil.deserializeObject(serializedAggregateFunction, userCodeLoader);
	} catch (Exception e) {
		log.error("Error while attempting to deserialize user AggregateFunction.");
		return FutureUtils.completedExceptionally(e);
	}

	Object accumulator = accumulators.get(aggregateName);
	if (null == accumulator) {
		accumulator = aggregateFunction.createAccumulator();
	}
	accumulator = aggregateFunction.add(aggregand, accumulator);
	accumulators.put(aggregateName, accumulator);
	return CompletableFuture.completedFuture(aggregateFunction.getResult(accumulator));
}
 

private <M extends MessageHeaders<EmptyRequestBody, P, U>, U extends MessageParameters, P extends ResponseBody> CompletableFuture<P>
		sendRequest(M messageHeaders, U messageParameters) {
	try {
		return restClient.sendRequest(
			restAddress.getHost(),
			restAddress.getPort(),
			messageHeaders,
			messageParameters,
			EmptyRequestBody.getInstance());
	} catch (IOException e) {
		return FutureUtils.completedExceptionally(e);
	}
}
 

/**
 * Called to shut down the TaskManager. The method closes all TaskManager services.
 */
@Override
public CompletableFuture<Void> onStop() {
	log.info("Stopping TaskExecutor {}.", getAddress());

	Throwable throwable = null;

	if (resourceManagerConnection != null) {
		resourceManagerConnection.close();
	}

	for (JobManagerConnection jobManagerConnection : jobManagerConnections.values()) {
		try {
			disassociateFromJobManager(jobManagerConnection, new FlinkException("The TaskExecutor is shutting down."));
		} catch (Throwable t) {
			throwable = ExceptionUtils.firstOrSuppressed(t, throwable);
		}
	}

	try {
		stopTaskExecutorServices();
	} catch (Exception e) {
		throwable = ExceptionUtils.firstOrSuppressed(e, throwable);
	}

	if (throwable != null) {
		return FutureUtils.completedExceptionally(new FlinkException("Error while shutting the TaskExecutor down.", throwable));
	} else {
		log.info("Stopped TaskExecutor {}.", getAddress());
		return CompletableFuture.completedFuture(null);
	}
}
 

@Override
public CompletableFuture<LogicalSlot> apply(SlotRequestId slotRequestId) {
	if (slotIterator.hasNext()) {
		return slotIterator.next();
	} else {
		return FutureUtils.completedExceptionally(new FlinkException("No more slots available."));
	}
}
 

@Override
public void handleError(Exception exception) {
	synchronized (lock) {
		if (!(leaderSessionFuture.isCompletedExceptionally() || leaderSessionFuture.completeExceptionally(exception))) {
			leaderSessionFuture = FutureUtils.completedExceptionally(exception);
		}
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendOperatorEventToTask(
		ExecutionAttemptID task,
		OperatorID operator,
		SerializedValue<OperatorEvent> evt) {
	return FutureUtils.completedExceptionally(new TestException());
}
 

@Override
public CompletableFuture<LogicalSlot> allocateSlot(
		SlotRequestId slotRequestId,
		ScheduledUnit task,
		SlotProfile slotProfile,
		Time allocationTimeout) {
	final SlotContext slot;

	synchronized (lock) {
		if (slots.isEmpty()) {
			slot = null;
		} else {
			slot = slots.removeFirst();
		}
		if (slot != null) {
			TestingLogicalSlot result = new TestingLogicalSlotBuilder()
				.setTaskManagerLocation(slot.getTaskManagerLocation())
				.setTaskManagerGateway(slot.getTaskManagerGateway())
				.setSlotNumber(slot.getPhysicalSlotNumber())
				.setAllocationId(slot.getAllocationId())
				.setSlotRequestId(slotRequestId)
				.setSlotSharingGroupId(task.getSlotSharingGroupId())
				.setSlotOwner(this)
				.createTestingLogicalSlot();
			allocatedSlots.put(slotRequestId, slot);
			return CompletableFuture.completedFuture(result);
		} else {
			return FutureUtils.completedExceptionally(new NoResourceAvailableException());
		}
	}
}
 

private CompletableFuture<Void> stopServices() {
	try {
		leaderElectionService.stop();
	} catch (Exception e) {
		return FutureUtils.completedExceptionally(e);
	}

	return FutureUtils.completedVoidFuture();
}
 

@Override
public CompletableFuture<Acknowledge> triggerCheckpoint(
		ExecutionAttemptID executionAttemptID,
		long checkpointId,
		long checkpointTimestamp,
		CheckpointOptions checkpointOptions,
		boolean advanceToEndOfEventTime) {
	log.debug("Trigger checkpoint {}@{} for {}.", checkpointId, checkpointTimestamp, executionAttemptID);

	final CheckpointType checkpointType = checkpointOptions.getCheckpointType();
	if (advanceToEndOfEventTime && !(checkpointType.isSynchronous() && checkpointType.isSavepoint())) {
		throw new IllegalArgumentException("Only synchronous savepoints are allowed to advance the watermark to MAX.");
	}

	final Task task = taskSlotTable.getTask(executionAttemptID);

	if (task != null) {
		task.triggerCheckpointBarrier(checkpointId, checkpointTimestamp, checkpointOptions, advanceToEndOfEventTime);

		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		final String message = "TaskManager received a checkpoint request for unknown task " + executionAttemptID + '.';

		log.debug(message);
		return FutureUtils.completedExceptionally(new CheckpointException(message, CheckpointFailureReason.TASK_CHECKPOINT_FAILURE));
	}
}
 

@Override
public CompletableFuture<KvStateLocation> requestKvStateLocation(final JobID jobId, final String registrationName) {
	try {
		return CompletableFuture.completedFuture(schedulerNG.requestKvStateLocation(jobId, registrationName));
	} catch (UnknownKvStateLocation | FlinkJobNotFoundException e) {
		log.info("Error while request key-value state location", e);
		return FutureUtils.completedExceptionally(e);
	}
}
 

@Override
public void handleError(Exception exception) {
	synchronized (lock) {
		if (!(leaderSessionFuture.isCompletedExceptionally() || leaderSessionFuture.completeExceptionally(exception))) {
			leaderSessionFuture = FutureUtils.completedExceptionally(exception);
		}
	}
}
 

/**
 * Lookup the {@link KvStateLocation} for the given job and queryable state name.
 *
 * <p>The job manager will be queried for the location only if forced or no
 * cached location can be found. There are no guarantees about
 *
 * @param jobId              JobID the state instance belongs to.
 * @param queryableStateName Name under which the state instance has been published.
 * @param forceUpdate        Flag to indicate whether to force a update via the lookup service.
 * @return Future holding the KvStateLocation
 */
private CompletableFuture<KvStateLocation> getKvStateLookupInfo(
		final JobID jobId,
		final String queryableStateName,
		final boolean forceUpdate) {

	final Tuple2<JobID, String> cacheKey = new Tuple2<>(jobId, queryableStateName);
	final CompletableFuture<KvStateLocation> cachedFuture = lookupCache.get(cacheKey);

	if (!forceUpdate && cachedFuture != null && !cachedFuture.isCompletedExceptionally()) {
		LOG.debug("Retrieving location for state={} of job={} from the cache.", queryableStateName, jobId);
		return cachedFuture;
	}

	final KvStateLocationOracle kvStateLocationOracle = proxy.getKvStateLocationOracle(jobId);

	if (kvStateLocationOracle != null) {
		LOG.debug("Retrieving location for state={} of job={} from the key-value state location oracle.", queryableStateName, jobId);
		final CompletableFuture<KvStateLocation> location = new CompletableFuture<>();
		lookupCache.put(cacheKey, location);

		kvStateLocationOracle
			.requestKvStateLocation(jobId, queryableStateName)
			.whenComplete(
				(KvStateLocation kvStateLocation, Throwable throwable) -> {
					if (throwable != null) {
						if (ExceptionUtils.stripCompletionException(throwable) instanceof FlinkJobNotFoundException) {
							// if the jobId was wrong, remove the entry from the cache.
							lookupCache.remove(cacheKey);
						}
						location.completeExceptionally(throwable);
					} else {
						location.complete(kvStateLocation);
					}
				});

		return location;
	} else {
		return FutureUtils.completedExceptionally(
			new UnknownLocationException(
					"Could not retrieve location of state=" + queryableStateName + " of job=" + jobId +
							". Potential reasons are: i) the state is not ready, or ii) the job does not exist."));
	}
}
 

@Override
public CompletableFuture<Acknowledge> requestSlot(
	final SlotID slotId,
	final JobID jobId,
	final AllocationID allocationId,
	final ResourceProfile resourceProfile,
	final String targetAddress,
	final ResourceManagerId resourceManagerId,
	final Time timeout) {
	// TODO: Filter invalid requests from the resource manager by using the instance/registration Id

	log.info("Receive slot request {} for job {} from resource manager with leader id {}.",
		allocationId, jobId, resourceManagerId);

	if (!isConnectedToResourceManager(resourceManagerId)) {
		final String message = String.format("TaskManager is not connected to the resource manager %s.", resourceManagerId);
		log.debug(message);
		return FutureUtils.completedExceptionally(new TaskManagerException(message));
	}

	try {
		allocateSlot(
			slotId,
			jobId,
			allocationId,
			resourceProfile);
	} catch (SlotAllocationException sae) {
		return FutureUtils.completedExceptionally(sae);
	}

	final JobTable.Job job;

	try {
		job = jobTable.getOrCreateJob(jobId, () -> registerNewJobAndCreateServices(jobId, targetAddress));
	} catch (Exception e) {
		// free the allocated slot
		try {
			taskSlotTable.freeSlot(allocationId);
		} catch (SlotNotFoundException slotNotFoundException) {
			// slot no longer existent, this should actually never happen, because we've
			// just allocated the slot. So let's fail hard in this case!
			onFatalError(slotNotFoundException);
		}

		// release local state under the allocation id.
		localStateStoresManager.releaseLocalStateForAllocationId(allocationId);

		// sanity check
		if (!taskSlotTable.isSlotFree(slotId.getSlotNumber())) {
			onFatalError(new Exception("Could not free slot " + slotId));
		}

		return FutureUtils.completedExceptionally(new SlotAllocationException("Could not create new job.", e));
	}

	if (job.isConnected()) {
		offerSlotsToJobManager(jobId);
	}

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

private CompletableFuture<LogicalSlot> allocateSharedSlot(
	SlotRequestId slotRequestId,
	ScheduledUnit scheduledUnit,
	SlotProfile slotProfile,
	boolean allowQueuedScheduling,
	@Nullable Time allocationTimeout) {
	// allocate slot with slot sharing
	final SlotSharingManager multiTaskSlotManager = slotSharingManagers.computeIfAbsent(
		scheduledUnit.getSlotSharingGroupId(),
		id -> new SlotSharingManager(
			id,
			slotPool,
			this));

	final SlotSharingManager.MultiTaskSlotLocality multiTaskSlotLocality;
	try {
		if (scheduledUnit.getCoLocationConstraint() != null) {
			multiTaskSlotLocality = allocateCoLocatedMultiTaskSlot(
				scheduledUnit.getCoLocationConstraint(),
				multiTaskSlotManager,
				slotProfile,
				allowQueuedScheduling,
				allocationTimeout);
		} else {
			multiTaskSlotLocality = allocateMultiTaskSlot(
				scheduledUnit.getJobVertexId(),
				multiTaskSlotManager,
				slotProfile,
				allowQueuedScheduling,
				allocationTimeout);
		}
	} catch (NoResourceAvailableException noResourceException) {
		return FutureUtils.completedExceptionally(noResourceException);
	}

	// sanity check
	Preconditions.checkState(!multiTaskSlotLocality.getMultiTaskSlot().contains(scheduledUnit.getJobVertexId()));

	final SlotSharingManager.SingleTaskSlot leaf = multiTaskSlotLocality.getMultiTaskSlot().allocateSingleTaskSlot(
		slotRequestId,
		slotProfile.getResourceProfile(),
		scheduledUnit.getJobVertexId(),
		multiTaskSlotLocality.getLocality());
	return leaf.getLogicalSlotFuture();
}