org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway Java Examples

private static ExecutionGraph createExecutionGraph(TaskManagerGateway gateway, int parallelism) throws Exception {
	final JobID jobId = new JobID();

	final JobVertex vertex = new JobVertex("vertex");
	vertex.setInvokableClass(NoOpInvokable.class);
	vertex.setParallelism(parallelism);

	final SlotProvider slotProvider = new SimpleSlotProvider(jobId, parallelism, gateway);

	ExecutionGraph simpleTestGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		jobId,
		slotProvider,
		new FixedDelayRestartStrategy(0, 0),
		vertex);
	simpleTestGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());
	return simpleTestGraph;
}
 

/**
 * Update the partition infos on the assigned resource.
 *
 * @param partitionInfos for the remote task
 */
private void sendUpdatePartitionInfoRpcCall(
		final Iterable<PartitionInfo> partitionInfos) {

	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
		final TaskManagerLocation taskManagerLocation = slot.getTaskManagerLocation();

		CompletableFuture<Acknowledge> updatePartitionsResultFuture = taskManagerGateway.updatePartitions(attemptId, partitionInfos, rpcTimeout);

		updatePartitionsResultFuture.whenCompleteAsync(
			(ack, failure) -> {
				// fail if there was a failure
				if (failure != null) {
					fail(new IllegalStateException("Update task on TaskManager " + taskManagerLocation +
						" failed due to:", failure));
				}
			}, getVertex().getExecutionGraph().getJobMasterMainThreadExecutor());
	}
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));



	LogicalSlot slot = new SingleLogicalSlot(
		new SlotRequestId(),
		slotContext,
		null,
		Locality.LOCAL,
		mock(SlotOwner.class));

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(slot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

/**
 * This method sends a CancelTask message to the instance of the assigned slot.
 *
 * <p>The sending is tried up to NUM_CANCEL_CALL_TRIES times.
 */
private void sendCancelRpcCall(int numberRetries) {
	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
		final ComponentMainThreadExecutor jobMasterMainThreadExecutor =
			getVertex().getExecutionGraph().getJobMasterMainThreadExecutor();

		CompletableFuture<Acknowledge> cancelResultFuture = FutureUtils.retry(
			() -> taskManagerGateway.cancelTask(attemptId, rpcTimeout),
			numberRetries,
			jobMasterMainThreadExecutor);

		cancelResultFuture.whenComplete(
			(ack, failure) -> {
				if (failure != null) {
					fail(new Exception("Task could not be canceled.", failure));
				}
			});
	}
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));

	SimpleSlot simpleSlot = new SimpleSlot(slot, mock(SlotOwner.class), 0);

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(simpleSlot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

/**
 * This method sends a CancelTask message to the instance of the assigned slot.
 *
 * <p>The sending is tried up to NUM_CANCEL_CALL_TRIES times.
 */
private void sendCancelRpcCall(int numberRetries) {
	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
		final ComponentMainThreadExecutor jobMasterMainThreadExecutor =
			getVertex().getExecutionGraph().getJobMasterMainThreadExecutor();

		CompletableFuture<Acknowledge> cancelResultFuture = FutureUtils.retry(
			() -> taskManagerGateway.cancelTask(attemptId, rpcTimeout),
			numberRetries,
			jobMasterMainThreadExecutor);

		cancelResultFuture.whenComplete(
			(ack, failure) -> {
				if (failure != null) {
					fail(new Exception("Task could not be canceled.", failure));
				}
			});
	}
}
 

TestingLogicalSlot(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		int slotNumber,
		AllocationID allocationId,
		SlotRequestId slotRequestId,
		SlotSharingGroupId slotSharingGroupId,
		boolean automaticallyCompleteReleaseFuture,
		SlotOwner slotOwner) {

	this.taskManagerLocation = Preconditions.checkNotNull(taskManagerLocation);
	this.taskManagerGateway = Preconditions.checkNotNull(taskManagerGateway);
	this.payloadReference = new AtomicReference<>();
	this.slotNumber = slotNumber;
	this.allocationId = Preconditions.checkNotNull(allocationId);
	this.slotRequestId = Preconditions.checkNotNull(slotRequestId);
	this.slotSharingGroupId = Preconditions.checkNotNull(slotSharingGroupId);
	this.releaseFuture = new CompletableFuture<>();
	this.automaticallyCompleteReleaseFuture = automaticallyCompleteReleaseFuture;
	this.slotOwner = slotOwner;
}
 

public SimpleSlotProvider(JobID jobId, int numSlots, TaskManagerGateway taskManagerGateway) {
	checkNotNull(jobId, "jobId");
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

private void triggerCheckpointHelper(long checkpointId, long timestamp, CheckpointOptions checkpointOptions, boolean advanceToEndOfEventTime) {

		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 LogicalSlot slot = assignedResource;

		if (slot != null) {
			final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();

			taskManagerGateway.triggerCheckpoint(attemptId, getVertex().getJobId(), checkpointId, timestamp, checkpointOptions, advanceToEndOfEventTime);
		} else {
			LOG.debug("The execution has no slot assigned. This indicates that the execution is no longer running.");
		}
	}
 

/**
 * Update the partition infos on the assigned resource.
 *
 * @param partitionInfos for the remote task
 */
private void sendUpdatePartitionInfoRpcCall(
		final Iterable<PartitionInfo> partitionInfos) {

	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
		final TaskManagerLocation taskManagerLocation = slot.getTaskManagerLocation();

		CompletableFuture<Acknowledge> updatePartitionsResultFuture = taskManagerGateway.updatePartitions(attemptId, partitionInfos, rpcTimeout);

		updatePartitionsResultFuture.whenCompleteAsync(
			(ack, failure) -> {
				// fail if there was a failure
				if (failure != null) {
					fail(new IllegalStateException("Update to task [" + getVertexWithAttempt() +
						"] on TaskManager " + taskManagerLocation + " failed", failure));
				}
			}, getVertex().getExecutionGraph().getJobMasterMainThreadExecutor());
	}
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private void triggerCheckpointHelper(long checkpointId, long timestamp, CheckpointOptions checkpointOptions, boolean advanceToEndOfEventTime) {

		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 LogicalSlot slot = assignedResource;

		if (slot != null) {
			final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();

			taskManagerGateway.triggerCheckpoint(attemptId, getVertex().getJobId(), checkpointId, timestamp, checkpointOptions, advanceToEndOfEventTime);
		} else {
			LOG.debug("The execution has no slot assigned. This indicates that the execution is no longer running.");
		}
	}
 

/**
 * Creates an execution graph containing the given vertices and the given restart strategy.
 */
public static ExecutionGraph createSimpleTestGraph(
		JobID jid,
		TaskManagerGateway taskManagerGateway,
		RestartStrategy restartStrategy,
		JobVertex... vertices) throws Exception {

	int numSlotsNeeded = 0;
	for (JobVertex vertex : vertices) {
		numSlotsNeeded += vertex.getParallelism();
	}

	SlotProvider slotProvider = new SimpleSlotProvider(jid, numSlotsNeeded, taskManagerGateway);

	return createSimpleTestGraph(jid, slotProvider, restartStrategy, vertices);
}
 

/**
 * Creates an execution graph containing the given vertices and the given restart strategy.
 */
public static ExecutionGraph createSimpleTestGraph(
		JobID jid,
		TaskManagerGateway taskManagerGateway,
		RestartStrategy restartStrategy,
		JobVertex... vertices) throws Exception {

	int numSlotsNeeded = 0;
	for (JobVertex vertex : vertices) {
		numSlotsNeeded += vertex.getParallelism();
	}

	SlotProvider slotProvider = new SimpleSlotProvider(jid, numSlotsNeeded, taskManagerGateway);

	return createSimpleTestGraph(jid, slotProvider, restartStrategy, vertices);
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

/**
 * Constructs an instance reflecting a registered TaskManager.
 *
 * @param taskManagerGateway The actor gateway to communicate with the remote instance
 * @param location The remote connection where the task manager receives requests.
 * @param id The id under which the taskManager is registered.
 * @param resources The resources available on the machine.
 * @param numberOfSlots The number of task slots offered by this taskManager.
 */
public Instance(
		TaskManagerGateway taskManagerGateway,
		TaskManagerLocation location,
		InstanceID id,
		HardwareDescription resources,
		int numberOfSlots) {
	this.taskManagerGateway = Preconditions.checkNotNull(taskManagerGateway);
	this.location = Preconditions.checkNotNull(location);
	this.instanceId = Preconditions.checkNotNull(id);
	this.resources = Preconditions.checkNotNull(resources);
	this.numberOfSlots = numberOfSlots;

	this.availableSlots = new ArrayDeque<>(numberOfSlots);
	for (int i = 0; i < numberOfSlots; i++) {
		this.availableSlots.add(i);
	}
}
 

public SimpleSlotProvider(int numSlots, TaskManagerGateway taskManagerGateway) {
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway,
			ResourceProfile.ANY);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

/**
 * Creates a new simple slot that belongs to the given shared slot and
 * is identified by the given ID.
 *
 * @param owner The component from which this slot is allocated.
 * @param location The location info of the TaskManager where the slot was allocated from
 * @param slotNumber The number of the simple slot in its parent shared slot.
 * @param taskManagerGateway to communicate with the associated task manager.
 * @param parent The parent shared slot.
 * @param groupID The ID that identifies the group that the slot belongs to.
 */
public SimpleSlot(
		SlotOwner owner,
		TaskManagerLocation location,
		int slotNumber,
		TaskManagerGateway taskManagerGateway,
		@Nullable SharedSlot parent,
		@Nullable AbstractID groupID) {

	super(
		parent != null ?
			parent.getSlotContext() :
			new SimpleSlotContext(
				NO_ALLOCATION_ID,
				location,
				slotNumber,
				taskManagerGateway),
		owner,
		slotNumber,
		parent,
		groupID);
}
 

/**
 * Request a stack trace sample from the task of this execution.
 *
 * @param sampleId of the stack trace sample
 * @param numSamples the sample should contain
 * @param delayBetweenSamples to wait
 * @param maxStackTraceDepth of the samples
 * @param timeout until the request times out
 * @return Future stack trace sample response
 */
public CompletableFuture<StackTraceSampleResponse> requestStackTraceSample(
		int sampleId,
		int numSamples,
		Time delayBetweenSamples,
		int maxStackTraceDepth,
		Time timeout) {

	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();

		return taskManagerGateway.requestStackTraceSample(
			attemptId,
			sampleId,
			numSamples,
			delayBetweenSamples,
			maxStackTraceDepth,
			timeout);
	} else {
		return FutureUtils.completedExceptionally(new Exception("The execution has no slot assigned."));
	}
}
 

private static ExecutionGraph createExecutionGraph(TaskManagerGateway gateway, int parallelism) throws Exception {
	final JobID jobId = new JobID();

	final JobVertex vertex = new JobVertex("vertex");
	vertex.setInvokableClass(NoOpInvokable.class);
	vertex.setParallelism(parallelism);

	final SlotProvider slotProvider = new SimpleSlotProvider(jobId, parallelism, gateway);

	ExecutionGraph simpleTestGraph = ExecutionGraphTestUtils.createSimpleTestGraph(
		jobId,
		slotProvider,
		new FixedDelayRestartStrategy(0, 0),
		vertex);
	simpleTestGraph.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());
	return simpleTestGraph;
}
 

/**
 * Update the partition infos on the assigned resource.
 *
 * @param partitionInfos for the remote task
 */
private void sendUpdatePartitionInfoRpcCall(
		final Iterable<PartitionInfo> partitionInfos) {

	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
		final TaskManagerLocation taskManagerLocation = slot.getTaskManagerLocation();

		CompletableFuture<Acknowledge> updatePartitionsResultFuture = taskManagerGateway.updatePartitions(attemptId, partitionInfos, rpcTimeout);

		updatePartitionsResultFuture.whenCompleteAsync(
			(ack, failure) -> {
				// fail if there was a failure
				if (failure != null) {
					fail(new IllegalStateException("Update task on TaskManager " + taskManagerLocation +
						" failed due to:", failure));
				}
			}, getVertex().getExecutionGraph().getJobMasterMainThreadExecutor());
	}
}
 

private void sendReleaseIntermediateResultPartitionsRpcCall() {
	LOG.info("Discarding the results produced by task execution {}.", attemptId);
	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();

		final ShuffleMaster<?> shuffleMaster = getVertex().getExecutionGraph().getShuffleMaster();

		Collection<ResultPartitionID> partitionIds = producedPartitions.values().stream()
			.filter(resultPartitionDeploymentDescriptor -> resultPartitionDeploymentDescriptor.getPartitionType().isPipelined())
			.map(ResultPartitionDeploymentDescriptor::getShuffleDescriptor)
			.peek(shuffleMaster::releasePartitionExternally)
			.map(ShuffleDescriptor::getResultPartitionID)
			.collect(Collectors.toList());

		if (!partitionIds.isEmpty()) {
			// TODO For some tests this could be a problem when querying too early if all resources were released
			taskManagerGateway.releasePartitions(getVertex().getJobId(), partitionIds);
		}
	}
}
 

@Override
public Collection<SlotOffer> offerSlots(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		Collection<SlotOffer> offers) {

	ArrayList<SlotOffer> result = new ArrayList<>(offers.size());

	for (SlotOffer offer : offers) {
		if (offerSlot(
			taskManagerLocation,
			taskManagerGateway,
			offer)) {

			result.add(offer);
		}
	}

	return result;
}
 

@Override
public Collection<SlotOffer> offerSlots(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		Collection<SlotOffer> offers) {

	ArrayList<SlotOffer> result = new ArrayList<>(offers.size());

	for (SlotOffer offer : offers) {
		if (offerSlot(
			taskManagerLocation,
			taskManagerGateway,
			offer)) {

			result.add(offer);
		}
	}

	return result;
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

private static Scheduler createSchedulerWithSlots(
	int numSlots, SlotPool slotPool, TaskManagerLocation taskManagerLocation) throws Exception {

	final TaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();
	setupSlotPool(slotPool);
	Scheduler scheduler = new SchedulerImpl(LocationPreferenceSlotSelectionStrategy.INSTANCE, slotPool);
	scheduler.start(mainThreadExecutor);
	slotPool.registerTaskManager(taskManagerLocation.getResourceID());

	final List<SlotOffer> slotOffers = new ArrayList<>(NUM_TASKS);
	for (int i = 0; i < numSlots; i++) {
		final AllocationID allocationId = new AllocationID();
		final SlotOffer slotOffer = new SlotOffer(allocationId, 0, ResourceProfile.ANY);
		slotOffers.add(slotOffer);
	}

	slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

	return scheduler;
}
 

private SimpleSlot createSlot(TaskManagerGateway taskManager, JobID jobId, SlotOwner slotOwner) {
	TaskManagerLocation location = new TaskManagerLocation(
			ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManager);

	return new SimpleSlot(slot, slotOwner, 0);
}
 

/**
 * Notify the task of this execution about a completed checkpoint.
 *
 * @param checkpointId of the completed checkpoint
 * @param timestamp of the completed checkpoint
 */
public void notifyCheckpointComplete(long checkpointId, long timestamp) {
	final LogicalSlot slot = assignedResource;

	if (slot != null) {
		final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();

		taskManagerGateway.notifyCheckpointComplete(attemptId, getVertex().getJobId(), checkpointId, timestamp);
	} else {
		LOG.debug("The execution has no slot assigned. This indicates that the execution is " +
			"no longer running.");
	}
}
 

public SimpleSlotContext(
		AllocationID allocationId,
		TaskManagerLocation taskManagerLocation,
		int physicalSlotNumber,
		TaskManagerGateway taskManagerGateway) {
	this(allocationId, taskManagerLocation, physicalSlotNumber, taskManagerGateway, ResourceProfile.UNKNOWN);
}