Java Code Examples for org.apache.flink.runtime.executiongraph.ExecutionAttemptID

@Override
public void triggerCheckpoint(
		ExecutionAttemptID executionAttemptID,
		JobID jobId,
		long checkpointId,
		long timestamp,
		CheckpointOptions checkpointOptions,
		boolean advanceToEndOfEventTime) {

	checkpointConsumer.accept(
		executionAttemptID,
		jobId,
		checkpointId,
		timestamp,
		checkpointOptions,
		advanceToEndOfEventTime);
}
 

/** Tests that collecting for a cancelled sample throws no Exception. */
@Test
public void testCollectStackTraceForCanceledSample() throws Exception {
	ExecutionVertex[] vertices = new ExecutionVertex[] {
			mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true),
	};

	CompletableFuture<StackTraceSample> sampleFuture = coord.triggerStackTraceSample(
			vertices, 1, Time.milliseconds(100L), 0);

	Assert.assertFalse(sampleFuture.isDone());

	coord.cancelStackTraceSample(0, null);

	Assert.assertTrue(sampleFuture.isDone());

	// Verify no error on late collect
	ExecutionAttemptID executionId = vertices[0].getCurrentExecutionAttempt().getAttemptId();
	coord.collectStackTraces(0, executionId, new ArrayList<StackTraceElement[]>());
}
 

static ExecutionVertex mockExecutionVertex(
	ExecutionAttemptID attemptID,
	JobVertexID jobVertexID,
	List<OperatorID> jobVertexIDs,
	int parallelism,
	int maxParallelism,
	ExecutionState state,
	ExecutionState ... successiveStates) {

	return mockExecutionVertex(
		attemptID,
		jobVertexID,
		jobVertexIDs,
		null,
		parallelism,
		maxParallelism,
		state,
		successiveStates);
}
 

private PendingCheckpoint createPendingCheckpoint(CheckpointProperties props, Executor executor) throws IOException {

		final Path checkpointDir = new Path(tmpFolder.newFolder().toURI());
		final FsCheckpointStorageLocation location = new FsCheckpointStorageLocation(
				LocalFileSystem.getSharedInstance(),
				checkpointDir, checkpointDir, checkpointDir,
				CheckpointStorageLocationReference.getDefault(),
				1024,
				4096);

		final Map<ExecutionAttemptID, ExecutionVertex> ackTasks = new HashMap<>(ACK_TASKS);

		return new PendingCheckpoint(
			new JobID(),
			0,
			1,
			ackTasks,
			props,
			location,
			executor);
	}
 

@Test
public void testEqualsHashCode() {
	try {
		final JobID jid = new JobID();
		final ExecutionAttemptID executionId = new ExecutionAttemptID();
		final ExecutionState state = ExecutionState.RUNNING;
		final Throwable error = new RuntimeException("some test error message");
		
		TaskExecutionState s1 = new TaskExecutionState(jid, executionId, state, error);
		TaskExecutionState s2 = new TaskExecutionState(jid, executionId, state, error);
		
		assertEquals(s1.hashCode(), s2.hashCode());
		assertEquals(s1, s2);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testEqualsHashCode() {
	try {
		final JobID jid = new JobID();
		final ExecutionAttemptID executionId = new ExecutionAttemptID();
		final ExecutionState state = ExecutionState.RUNNING;
		final Throwable error = new RuntimeException("some test error message");
		
		TaskExecutionState s1 = new TaskExecutionState(jid, executionId, state, error);
		TaskExecutionState s2 = new TaskExecutionState(jid, executionId, state, error);
		
		assertEquals(s1.hashCode(), s2.hashCode());
		assertEquals(s1, s2);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests that when partition consumable notified will start available {@link ResultPartitionType#PIPELINED} downstream vertices.
 */
@Test
public void testPipelinedPartitionConsumable() {
	final TestingSchedulingTopology testingSchedulingTopology = new TestingSchedulingTopology();

	final List<TestingSchedulingExecutionVertex> producers = testingSchedulingTopology.addExecutionVertices()
		.withParallelism(2).finish();
	final List<TestingSchedulingExecutionVertex> consumers = testingSchedulingTopology.addExecutionVertices()
		.withParallelism(2).finish();
	testingSchedulingTopology.connectAllToAll(producers, consumers).withResultPartitionType(PIPELINED).finish();

	final LazyFromSourcesSchedulingStrategy schedulingStrategy = startScheduling(testingSchedulingTopology);

	final TestingSchedulingExecutionVertex producer1 = producers.get(0);
	final SchedulingResultPartition partition1 = producer1.getProducedResultPartitions().iterator().next();

	schedulingStrategy.onExecutionStateChange(producer1.getId(), ExecutionState.RUNNING);
	schedulingStrategy.onPartitionConsumable(producer1.getId(), new ResultPartitionID(partition1.getId(), new ExecutionAttemptID()));

	assertThat(testingSchedulerOperation, hasScheduledVertices(consumers));
}
 

@Override
public CompletableFuture<StackTraceSampleResponse> requestStackTraceSample(
		ExecutionAttemptID executionAttemptID,
		int sampleId,
		int numSamples,
		Time delayBetweenSamples,
		int maxStackTraceDepth,
		Time timeout) {

	return taskExecutorGateway.requestStackTraceSample(
		executionAttemptID,
		sampleId,
		numSamples,
		delayBetweenSamples,
		maxStackTraceDepth,
		timeout);
}
 

TestingTaskExecutorGateway(
		String address,
		String hostname,
		BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer,
		BiConsumer<JobID, Throwable> disconnectJobManagerConsumer,
		BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer,
		Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction,
		BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>> freeSlotFunction,
		Consumer<ResourceID> heartbeatResourceManagerConsumer,
		Consumer<Exception> disconnectResourceManagerConsumer,
		Function<ExecutionAttemptID, CompletableFuture<Acknowledge>> cancelTaskFunction,
		Supplier<CompletableFuture<Boolean>> canBeReleasedSupplier,
		BiConsumer<JobID, Collection<ResultPartitionID>> releasePartitionsConsumer) {
	this.address = Preconditions.checkNotNull(address);
	this.hostname = Preconditions.checkNotNull(hostname);
	this.heartbeatJobManagerConsumer = Preconditions.checkNotNull(heartbeatJobManagerConsumer);
	this.disconnectJobManagerConsumer = Preconditions.checkNotNull(disconnectJobManagerConsumer);
	this.submitTaskConsumer = Preconditions.checkNotNull(submitTaskConsumer);
	this.requestSlotFunction = Preconditions.checkNotNull(requestSlotFunction);
	this.freeSlotFunction = Preconditions.checkNotNull(freeSlotFunction);
	this.heartbeatResourceManagerConsumer = heartbeatResourceManagerConsumer;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.cancelTaskFunction = cancelTaskFunction;
	this.canBeReleasedSupplier = canBeReleasedSupplier;
	this.releasePartitionsConsumer = releasePartitionsConsumer;
}
 

TestingTaskExecutorGateway(
		String address,
		String hostname,
		BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer,
		BiConsumer<JobID, Throwable> disconnectJobManagerConsumer,
		BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer,
		Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction,
		BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>> freeSlotFunction,
		Consumer<ResourceID> heartbeatResourceManagerConsumer,
		Consumer<Exception> disconnectResourceManagerConsumer,
		Function<ExecutionAttemptID, CompletableFuture<Acknowledge>> cancelTaskFunction,
		Supplier<Boolean> canBeReleasedSupplier) {
	this.address = Preconditions.checkNotNull(address);
	this.hostname = Preconditions.checkNotNull(hostname);
	this.heartbeatJobManagerConsumer = Preconditions.checkNotNull(heartbeatJobManagerConsumer);
	this.disconnectJobManagerConsumer = Preconditions.checkNotNull(disconnectJobManagerConsumer);
	this.submitTaskConsumer = Preconditions.checkNotNull(submitTaskConsumer);
	this.requestSlotFunction = Preconditions.checkNotNull(requestSlotFunction);
	this.freeSlotFunction = Preconditions.checkNotNull(freeSlotFunction);
	this.heartbeatResourceManagerConsumer = heartbeatResourceManagerConsumer;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.cancelTaskFunction = cancelTaskFunction;
	this.canBeReleasedSupplier = canBeReleasedSupplier;
}
 

/**
 * Tests that we can submit a task to the TaskManager given that we've allocated a slot there.
 */
@Test(timeout = 10000L)
public void testTaskSubmission() throws Exception {
	final ExecutionAttemptID eid = new ExecutionAttemptID();

	final TaskDeploymentDescriptor tdd = createTestTaskDeploymentDescriptor("test task", eid, TaskExecutorTest.TestInvokable.class);

	final CompletableFuture<Void> taskRunningFuture = new CompletableFuture<>();

	try (TaskSubmissionTestEnvironment env =
		new TaskSubmissionTestEnvironment.Builder(jobId)
			.setSlotSize(1)
			.addTaskManagerActionListener(eid, ExecutionState.RUNNING, taskRunningFuture)
			.build()) {
		TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
		TaskSlotTable taskSlotTable = env.getTaskSlotTable();

		taskSlotTable.allocateSlot(0, jobId, tdd.getAllocationId(), Time.seconds(60));
		tmGateway.submitTask(tdd, env.getJobMasterId(), timeout).get();

		taskRunningFuture.get();
	}
}
 

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

/**
 * Creates a stack trace sample.
 *
 * @param sampleId          ID of the sample.
 * @param startTime         Time stamp, when the sample was triggered.
 * @param endTime           Time stamp, when all stack traces were
 *                          collected at the JobManager.
 * @param stackTracesByTask Map of stack traces by execution ID.
 */
public StackTraceSample(
		int sampleId,
		long startTime,
		long endTime,
		Map<ExecutionAttemptID, List<StackTraceElement[]>> stackTracesByTask) {

	checkArgument(sampleId >= 0, "Negative sample ID");
	checkArgument(startTime >= 0, "Negative start time");
	checkArgument(endTime >= startTime, "End time before start time");

	this.sampleId = sampleId;
	this.startTime = startTime;
	this.endTime = endTime;
	this.stackTracesByTask = Collections.unmodifiableMap(stackTracesByTask);
}
 

@Override
public CompletableFuture<StackTraceSampleResponse> requestStackTraceSample(
		final ExecutionAttemptID executionAttemptId,
		final int sampleId,
		final int numSamples,
		final Time delayBetweenSamples,
		final int maxStackTraceDepth,
		final Time timeout) {

	final Task task = taskSlotTable.getTask(executionAttemptId);
	if (task == null) {
		return FutureUtils.completedExceptionally(
			new IllegalStateException(String.format("Cannot sample task %s. " +
				"Task is not known to the task manager.", executionAttemptId)));
	}

	final CompletableFuture<List<StackTraceElement[]>> stackTracesFuture = stackTraceSampleService.requestStackTraceSample(
		TaskStackTraceSampleableTaskAdapter.fromTask(task),
		numSamples,
		delayBetweenSamples,
		maxStackTraceDepth);

	return stackTracesFuture.thenApply(stackTraces ->
		new StackTraceSampleResponse(sampleId, executionAttemptId, stackTraces));
}
 

public void releaseJob(JobID jobId, ExecutionAttemptID executionId) {
	checkNotNull(jobId);

	synchronized (lock) {
		Set<ExecutionAttemptID> jobRefCounter = jobRefHolders.get(jobId);

		if (jobRefCounter == null || jobRefCounter.isEmpty()) {
			return;
		}

		jobRefCounter.remove(executionId);
		if (jobRefCounter.isEmpty()) {
			executorService.schedule(new DeleteProcess(jobId), cleanupInterval, TimeUnit.MILLISECONDS);
		}
	}
}
 

private void failPendingCheckpointDueToTaskFailure(
	final PendingCheckpoint pendingCheckpoint,
	final CheckpointFailureReason reason,
	final ExecutionAttemptID executionAttemptID) {

	failPendingCheckpointDueToTaskFailure(pendingCheckpoint, reason, null, executionAttemptID);
}
 

@Override
public void notifyTaskFailure(final ExecutionAttemptID attemptId, final Throwable t) {
	schedulerNg.updateTaskExecutionState(new TaskExecutionState(
		jobId,
		attemptId,
		ExecutionState.FAILED,
		t));
}
 

public AcknowledgeReport(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	long checkpointId,
	CheckpointMetrics checkpointMetrics,
	TaskStateSnapshot subtaskState) {

	super(jobID, executionAttemptID, checkpointId);
	this.checkpointMetrics = checkpointMetrics;
	this.subtaskState = subtaskState;
}
 

public StreamMockEnvironment(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	Configuration jobConfig,
	Configuration taskConfig,
	ExecutionConfig executionConfig,
	long offHeapMemorySize,
	MockInputSplitProvider inputSplitProvider,
	int bufferSize,
	TaskStateManager taskStateManager) {

	this.jobID = jobID;
	this.executionAttemptID = executionAttemptID;

	int subtaskIndex = 0;
	this.taskInfo = new TaskInfo(
		"", /* task name */
		1, /* num key groups / max parallelism */
		subtaskIndex, /* index of this subtask */
		1, /* num subtasks */
		0 /* attempt number */);
	this.jobConfiguration = jobConfig;
	this.taskConfiguration = taskConfig;
	this.inputs = new LinkedList<>();
	this.outputs = new LinkedList<ResultPartitionWriter>();
	this.memManager = MemoryManagerBuilder.newBuilder().setMemorySize(offHeapMemorySize).build();
	this.ioManager = new IOManagerAsync();
	this.taskStateManager = Preconditions.checkNotNull(taskStateManager);
	this.aggregateManager = new TestGlobalAggregateManager();
	this.inputSplitProvider = inputSplitProvider;
	this.bufferSize = bufferSize;

	this.executionConfig = executionConfig;
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, getExecutionId());

	KvStateRegistry registry = new KvStateRegistry();
	this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
}
 

/**
 * Tests that no minimum delay between savepoints is enforced.
 */
@Test
public void testMinDelayBetweenSavepoints() throws Exception {
	JobID jobId = new JobID();

	final ExecutionAttemptID attemptID1 = new ExecutionAttemptID();
	ExecutionVertex vertex1 = mockExecutionVertex(attemptID1);

	CheckpointCoordinatorConfiguration chkConfig = new CheckpointCoordinatorConfiguration(
		100000,
		200000,
		100000000L, // very long min delay => should not affect savepoints
		1,
		CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
		true,
		false,
		0);
	CheckpointCoordinator coord = new CheckpointCoordinator(
		jobId,
		chkConfig,
		new ExecutionVertex[] { vertex1 },
		new ExecutionVertex[] { vertex1 },
		new ExecutionVertex[] { vertex1 },
		new StandaloneCheckpointIDCounter(),
		new StandaloneCompletedCheckpointStore(2),
		new MemoryStateBackend(),
		Executors.directExecutor(),
		SharedStateRegistry.DEFAULT_FACTORY,
		failureManager);

	String savepointDir = tmpFolder.newFolder().getAbsolutePath();

	CompletableFuture<CompletedCheckpoint> savepoint0 = coord.triggerSavepoint(0, savepointDir);
	assertFalse("Did not trigger savepoint", savepoint0.isDone());

	CompletableFuture<CompletedCheckpoint> savepoint1 = coord.triggerSavepoint(1, savepointDir);
	assertFalse("Did not trigger savepoint", savepoint1.isDone());
}
 

/**
 * Collects stack traces of a task.
 *
 * @param sampleId    ID of the sample.
 * @param executionId ID of the sampled task.
 * @param stackTraces Stack traces of the sampled task.
 *
 * @throws IllegalStateException If unknown sample ID and not recently
 *                               finished or cancelled sample.
 */
public void collectStackTraces(
		int sampleId,
		ExecutionAttemptID executionId,
		List<StackTraceElement[]> stackTraces) {

	synchronized (lock) {
		if (isShutDown) {
			return;
		}

		if (LOG.isDebugEnabled()) {
			LOG.debug("Collecting stack trace sample {} of task {}", sampleId, executionId);
		}

		PendingStackTraceSample pending = pendingSamples.get(sampleId);

		if (pending != null) {
			pending.collectStackTraces(executionId, stackTraces);

			// Publish the sample
			if (pending.isComplete()) {
				pendingSamples.remove(sampleId);
				rememberRecentSampleId(sampleId);

				pending.completePromiseAndDiscard();
			}
		} else if (recentPendingSamples.contains(sampleId)) {
			if (LOG.isDebugEnabled()) {
				LOG.debug("Received late stack trace sample {} of task {}",
						sampleId, executionId);
			}
		} else {
			if (LOG.isDebugEnabled()) {
				LOG.debug("Unknown sample ID " + sampleId);
			}
		}
	}
}
 

@Override
public void run() {
	try {
		synchronized (lock) {

			Set<ExecutionAttemptID> jobRefs = jobRefHolders.get(jobID);
			if (jobRefs != null && jobRefs.isEmpty()) {
				// abort the copy
				for (Future<Path> fileFuture : entries.get(jobID).values()) {
					fileFuture.cancel(true);
				}

				//remove job specific entries in maps
				entries.remove(jobID);
				jobRefHolders.remove(jobID);

				// remove the job wide temp directories
				for (File storageDirectory : storageDirectories) {
					File tempDir = new File(storageDirectory, jobID.toString());
					FileUtils.deleteDirectory(tempDir);
				}
			}
		}
	} catch (IOException e) {
		LOG.error("Could not delete file from local file cache.", e);
	}
}
 

private PendingCheckpoint createPendingCheckpoint(
		CheckpointProperties props,
		Collection<OperatorID> operatorCoordinators,
		Collection<String> masterStateIdentifiers,
		Executor executor) throws IOException {

	final Path checkpointDir = new Path(tmpFolder.newFolder().toURI());
	final FsCheckpointStorageLocation location = new FsCheckpointStorageLocation(
			LocalFileSystem.getSharedInstance(),
			checkpointDir, checkpointDir, checkpointDir,
			CheckpointStorageLocationReference.getDefault(),
			1024,
			4096);

	final Map<ExecutionAttemptID, ExecutionVertex> ackTasks = new HashMap<>(ACK_TASKS);

	return new PendingCheckpoint(
		new JobID(),
		0,
		1,
		ackTasks,
		operatorCoordinators,
		masterStateIdentifiers,
		props,
		location,
		executor,
		new CompletableFuture<>());
}
 

@Override
public void acknowledgeCheckpoint(
	JobID jobID,
	ExecutionAttemptID executionAttemptID,
	long checkpointId,
	CheckpointMetrics checkpointMetrics,
	TaskStateSnapshot subtaskState) {

	throw new RuntimeException("Unexpected call.");
}
 

/** Tests that collecting for a unknown task fails. */
@Test(expected = IllegalArgumentException.class)
public void testCollectStackTraceForUnknownTask() throws Exception {
	ExecutionVertex[] vertices = new ExecutionVertex[] {
			mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true),
	};

	coord.triggerStackTraceSample(vertices, 1, Time.milliseconds(100L), 0);

	coord.collectStackTraces(0, new ExecutionAttemptID(), new ArrayList<StackTraceElement[]>());
}
 

@Override
public void notifyCheckpointComplete(
		ExecutionAttemptID executionAttemptID,
		JobID jobId,
		long checkpointId,
		long timestamp) {

	Preconditions.checkNotNull(executionAttemptID);
	Preconditions.checkNotNull(jobId);

	actorGateway.tell(new NotifyCheckpointComplete(jobId, executionAttemptID, checkpointId, timestamp));
}
 

private void unregisterTaskAndNotifyFinalState(
		final JobMasterGateway jobMasterGateway,
		final ExecutionAttemptID executionAttemptID) {

	Task task = taskSlotTable.removeTask(executionAttemptID);
	if (task != null) {
		if (!task.getExecutionState().isTerminal()) {
			try {
				task.failExternally(new IllegalStateException("Task is being remove from TaskManager."));
			} catch (Exception e) {
				log.error("Could not properly fail task.", e);
			}
		}

		log.info("Un-registering task and sending final execution state {} to JobManager for task {} {}.",
			task.getExecutionState(), task.getTaskInfo().getTaskName(), task.getExecutionId());

		AccumulatorSnapshot accumulatorSnapshot = task.getAccumulatorRegistry().getSnapshot();

		updateTaskExecutionState(
				jobMasterGateway,
				new TaskExecutionState(
					task.getJobID(),
					task.getExecutionId(),
					task.getExecutionState(),
					task.getFailureCause(),
					accumulatorSnapshot,
					task.getMetricGroup().getIOMetricGroup().createSnapshot()));
	} else {
		log.error("Cannot find task with ID {} to unregister.", executionAttemptID);
	}
}
 

public StackTraceSampleResponse(
		int sampleId,
		ExecutionAttemptID executionAttemptID,
		List<StackTraceElement[]> samples) {
	this.sampleId = sampleId;
	this.executionAttemptID = Preconditions.checkNotNull(executionAttemptID);
	this.samples = Preconditions.checkNotNull(samples);
}
 

private String retrieveTaskManagerLocation(ExecutionAttemptID executionAttemptID) {
	final Optional<Execution> currentExecution = Optional.ofNullable(executionGraph.getRegisteredExecutions().get(executionAttemptID));

	return currentExecution
		.map(Execution::getAssignedResourceLocation)
		.map(TaskManagerLocation::toString)
		.orElse("Unknown location");
}
 

public void unregisterTask(Task task) {
	LOG.debug("Unregister task {} from network environment (state: {}).",
			task.getTaskInfo().getTaskNameWithSubtasks(), task.getExecutionState());

	final ExecutionAttemptID executionId = task.getExecutionId();

	synchronized (lock) {
		if (isShutdown) {
			// no need to do anything when we are not operational
			return;
		}

		if (task.isCanceledOrFailed()) {
			resultPartitionManager.releasePartitionsProducedBy(executionId, task.getFailureCause());
		}

		for (ResultPartition partition : task.getProducedPartitions()) {
			taskEventDispatcher.unregisterPartition(partition.getPartitionId());
			partition.destroyBufferPool();
		}

		final SingleInputGate[] inputGates = task.getAllInputGates();

		if (inputGates != null) {
			for (SingleInputGate gate : inputGates) {
				try {
					if (gate != null) {
						gate.releaseAllResources();
					}
				}
				catch (IOException e) {
					LOG.error("Error during release of reader resources: " + e.getMessage(), e);
				}
			}
		}
	}
}