Java Code Examples for org.apache.flink.runtime.jobgraph.JobVertex#setSlotSharingGroup()

private JobGraph createJobGraph(int parallelism) {
	final JobVertex sender = new JobVertex("Sender");
	sender.setParallelism(parallelism);
	sender.setInvokableClass(TestingAbstractInvokables.Sender.class);

	final JobVertex receiver = new JobVertex("Blocking receiver");
	receiver.setParallelism(parallelism);
	receiver.setInvokableClass(BlockingOperator.class);
	BlockingOperator.reset();

	receiver.connectNewDataSetAsInput(sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setSlotSharingGroup(slotSharingGroup);

	return new JobGraph("Blocking test job with slot sharing", sender, receiver);
}
 

private JobGraph createJobGraph(int parallelism) {
	final JobVertex sender = new JobVertex("Sender");
	sender.setParallelism(parallelism);
	sender.setInvokableClass(TestingAbstractInvokables.Sender.class);

	final JobVertex receiver = new JobVertex("Receiver");
	receiver.setParallelism(parallelism);
	receiver.setInvokableClass(TestingAbstractInvokables.Receiver.class);

	// In order to make testCoLocationConstraintJobExecution fail, one needs to
	// remove the co-location constraint and the slot sharing groups, because then
	// the receivers will have to wait for the senders to finish and the slot
	// assignment order to the receivers is non-deterministic (depending on the
	// order in which the senders finish).
	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	receiver.setSlotSharingGroup(slotSharingGroup);
	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setStrictlyCoLocatedWith(sender);

	receiver.connectNewDataSetAsInput(sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final JobGraph jobGraph = new JobGraph(getClass().getSimpleName(), sender, receiver);

	return jobGraph;
}
 

private JobGraph createJobGraph(int parallelism) {
	final JobVertex sender = new JobVertex("Sender");
	sender.setParallelism(parallelism);
	sender.setInvokableClass(TestingAbstractInvokables.Sender.class);

	final JobVertex receiver = new JobVertex("Receiver");
	receiver.setParallelism(parallelism);
	receiver.setInvokableClass(TestingAbstractInvokables.Receiver.class);

	// In order to make testCoLocationConstraintJobExecution fail, one needs to
	// remove the co-location constraint and the slot sharing groups, because then
	// the receivers will have to wait for the senders to finish and the slot
	// assignment order to the receivers is non-deterministic (depending on the
	// order in which the senders finish).
	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	receiver.setSlotSharingGroup(slotSharingGroup);
	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setStrictlyCoLocatedWith(sender);

	receiver.connectNewDataSetAsInput(sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final JobGraph jobGraph = new JobGraph(getClass().getSimpleName(), sender, receiver);

	return jobGraph;
}
 

private JobGraph createJobGraph(int parallelism) {
	final JobVertex sender = new JobVertex("Sender");
	sender.setParallelism(parallelism);
	sender.setInvokableClass(TestingAbstractInvokables.Sender.class);

	final JobVertex receiver = new JobVertex("Blocking receiver");
	receiver.setParallelism(parallelism);
	receiver.setInvokableClass(BlockingOperator.class);
	BlockingOperator.reset();

	receiver.connectNewDataSetAsInput(sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setSlotSharingGroup(slotSharingGroup);

	return new JobGraph("Blocking test job with slot sharing", sender, receiver);
}
 

private static JobGraph createJobGraph() {
	final SlotSharingGroup group1 = new SlotSharingGroup();
	final SlotSharingGroup group2 = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(TestSourceInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(group1);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(TestSinkInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(group2);

	sink.connectNewDataSetAsInput(source, DistributionPattern.ALL_TO_ALL, ResultPartitionType.BLOCKING);

	final JobGraph jobGraph = new JobGraph(source, sink);
	jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES);

	return jobGraph;
}
 

@Nonnull
private JobGraph createJobGraph(long delay, int parallelism) throws IOException {
	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(OneTimeFailingInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(slotSharingGroup);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(slotSharingGroup);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
	JobGraph jobGraph = new JobGraph(source, sink);

	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, delay));
	jobGraph.setExecutionConfig(executionConfig);

	return jobGraph;
}
 

private JobGraph createJobGraph(int parallelism) {
	final JobVertex sender = new JobVertex("Sender");
	sender.setParallelism(parallelism);
	sender.setInvokableClass(TestingAbstractInvokables.Sender.class);

	final JobVertex receiver = new JobVertex("Receiver");
	receiver.setParallelism(parallelism);
	receiver.setInvokableClass(TestingAbstractInvokables.Receiver.class);

	// In order to make testCoLocationConstraintJobExecution fail, one needs to
	// remove the co-location constraint and the slot sharing groups, because then
	// the receivers will have to wait for the senders to finish and the slot
	// assignment order to the receivers is non-deterministic (depending on the
	// order in which the senders finish).
	final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	receiver.setSlotSharingGroup(slotSharingGroup);
	sender.setSlotSharingGroup(slotSharingGroup);
	receiver.setStrictlyCoLocatedWith(sender);

	receiver.connectNewDataSetAsInput(sender, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final JobGraph jobGraph = new JobGraph(getClass().getSimpleName(), sender, receiver);

	return jobGraph;
}
 

@Nonnull
private JobGraph createJobGraph(long delay, int parallelism) throws IOException {
	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(OneTimeFailingInvokable.class);
	source.setParallelism(parallelism);
	source.setSlotSharingGroup(slotSharingGroup);

	final JobVertex sink = new JobVertex("sink");
	sink.setInvokableClass(NoOpInvokable.class);
	sink.setParallelism(parallelism);
	sink.setSlotSharingGroup(slotSharingGroup);

	sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
	JobGraph jobGraph = new JobGraph(source, sink);

	jobGraph.setScheduleMode(ScheduleMode.EAGER);

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, delay));
	jobGraph.setExecutionConfig(executionConfig);

	return jobGraph;
}
 

/**
 * This test checks that are strictly co-located vertices are in the same failover region,
 * even through they are connected via a blocking pattern.
 * This is currently an assumption / limitation of the scheduler.
 */
@Test
public void testPipelinedOneToOneTopologyWithCoLocation() throws Exception {
	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(NoOpInvokable.class);
	source.setParallelism(10);

	final JobVertex target = new JobVertex("target");
	target.setInvokableClass(NoOpInvokable.class);
	target.setParallelism(10);

	target.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

	final SlotSharingGroup sharingGroup = new SlotSharingGroup();
	source.setSlotSharingGroup(sharingGroup);
	target.setSlotSharingGroup(sharingGroup);

	source.setStrictlyCoLocatedWith(target);

	final JobGraph jobGraph = new JobGraph("test job", source, target);
	final ExecutionGraph eg = createExecutionGraph(jobGraph);

	RestartPipelinedRegionStrategy failoverStrategy = (RestartPipelinedRegionStrategy) eg.getFailoverStrategy();
	FailoverRegion sourceRegion1 = failoverStrategy.getFailoverRegion(eg.getJobVertex(source.getID()).getTaskVertices()[0]);
	FailoverRegion sourceRegion2 = failoverStrategy.getFailoverRegion(eg.getJobVertex(source.getID()).getTaskVertices()[1]);
	FailoverRegion targetRegion1 = failoverStrategy.getFailoverRegion(eg.getJobVertex(target.getID()).getTaskVertices()[0]);
	FailoverRegion targetRegion2 = failoverStrategy.getFailoverRegion(eg.getJobVertex(target.getID()).getTaskVertices()[1]);

	// we use 'assertTrue' here rather than 'assertEquals' because we want to test
	// for referential equality, to be on the safe side
	assertTrue(sourceRegion1 == sourceRegion2);
	assertTrue(sourceRegion2 == targetRegion1);
	assertTrue(targetRegion1 == targetRegion2);
}
 

@Test
public void coLocationConstraintIsResetOnTaskRecovery() {
	final JobGraph jobGraph = nonParallelSourceSinkJobGraph();
	final JobVertex source = jobGraph.getVerticesSortedTopologicallyFromSources().get(0);
	final JobVertex sink = jobGraph.getVerticesSortedTopologicallyFromSources().get(1);

	final SlotSharingGroup ssg = new SlotSharingGroup();
	source.setSlotSharingGroup(ssg);
	sink.setSlotSharingGroup(ssg);
	sink.setStrictlyCoLocatedWith(source);

	final JobID jobId = jobGraph.getJobID();
	final DefaultScheduler scheduler = createSchedulerAndStartScheduling(jobGraph);

	final ExecutionVertex sourceVertex = scheduler.getExecutionVertex(new ExecutionVertexID(source.getID(), 0));
	final ExecutionAttemptID sourceAttemptId = sourceVertex.getCurrentExecutionAttempt().getAttemptId();
	final ExecutionVertex sinkVertex = scheduler.getExecutionVertex(new ExecutionVertexID(sink.getID(), 0));
	final ExecutionAttemptID sinkAttemptId = sinkVertex.getCurrentExecutionAttempt().getAttemptId();

	// init the location constraint manually because the testExecutionSlotAllocator does not do it
	sourceVertex.getLocationConstraint().setSlotRequestId(new SlotRequestId());
	assertThat(sourceVertex.getLocationConstraint().getSlotRequestId(), is(notNullValue()));

	final String exceptionMessage = "expected exception";
	scheduler.updateTaskExecutionState(new TaskExecutionState(jobId, sourceAttemptId, ExecutionState.FAILED, new RuntimeException(exceptionMessage)));
	scheduler.updateTaskExecutionState(new TaskExecutionState(jobId, sinkAttemptId, ExecutionState.CANCELED));
	taskRestartExecutor.triggerScheduledTasks();

	assertThat(sourceVertex.getLocationConstraint().getSlotRequestId(), is(nullValue()));
}
 

/**
 * This test checks that are strictly co-located vertices are in the same failover region,
 * even through they are connected via a blocking pattern.
 * This is currently an assumption / limitation of the scheduler.
 */
@Test
public void testBlockingAllToAllTopologyWithCoLocation() throws Exception {
	final JobVertex source = new JobVertex("source");
	source.setInvokableClass(NoOpInvokable.class);
	source.setParallelism(10);

	final JobVertex target = new JobVertex("target");
	target.setInvokableClass(NoOpInvokable.class);
	target.setParallelism(13);

	target.connectNewDataSetAsInput(source, DistributionPattern.ALL_TO_ALL, ResultPartitionType.BLOCKING);

	final SlotSharingGroup sharingGroup = new SlotSharingGroup();
	source.setSlotSharingGroup(sharingGroup);
	target.setSlotSharingGroup(sharingGroup);

	source.setStrictlyCoLocatedWith(target);

	final JobGraph jobGraph = new JobGraph("test job", source, target);
	final ExecutionGraph eg = createExecutionGraph(jobGraph);

	RestartPipelinedRegionStrategy failoverStrategy = (RestartPipelinedRegionStrategy) eg.getFailoverStrategy();
	FailoverRegion region1 = failoverStrategy.getFailoverRegion(eg.getJobVertex(source.getID()).getTaskVertices()[0]);
	FailoverRegion region2 = failoverStrategy.getFailoverRegion(eg.getJobVertex(target.getID()).getTaskVertices()[0]);

	// we use 'assertTrue' here rather than 'assertEquals' because we want to test
	// for referential equality, to be on the safe side
	assertTrue(region1 == region2);
}
 

@Test
public void testConstraintsAfterRestart() throws Exception {

	final long timeout = 5000L;

	//setting up
	testingSlotProvider.addTaskManager(NUM_TASKS);

	JobVertex groupVertex = ExecutionGraphTestUtils.createNoOpVertex(NUM_TASKS);
	JobVertex groupVertex2 = ExecutionGraphTestUtils.createNoOpVertex(NUM_TASKS);

	SlotSharingGroup sharingGroup = new SlotSharingGroup();
	groupVertex.setSlotSharingGroup(sharingGroup);
	groupVertex2.setSlotSharingGroup(sharingGroup);
	groupVertex.setStrictlyCoLocatedWith(groupVertex2);

	//initiate and schedule job
	final ExecutionGraph eg = TestingExecutionGraphBuilder
		.newBuilder()
		.setJobGraph(new JobGraph(groupVertex, groupVertex2))
		.setSlotProvider(testingSlotProvider)
		.setRestartStrategy(new TestRestartStrategy(1, false))
		.build();

	// enable the queued scheduling for the slot pool
	eg.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	assertEquals(JobStatus.CREATED, eg.getState());

	eg.scheduleForExecution();

	Predicate<AccessExecution> isDeploying = ExecutionGraphTestUtils.isInExecutionState(ExecutionState.DEPLOYING);
	ExecutionGraphTestUtils.waitForAllExecutionsPredicate(
		eg,
		isDeploying,
		timeout);

	assertEquals(JobStatus.RUNNING, eg.getState());

	//sanity checks
	validateConstraints(eg);

	eg.getAllExecutionVertices().iterator().next().fail(new FlinkException("Test exception"));

	assertEquals(JobStatus.FAILING, eg.getState());

	for (ExecutionVertex vertex : eg.getAllExecutionVertices()) {
		vertex.getCurrentExecutionAttempt().completeCancelling();
	}

	// wait until we have restarted
	ExecutionGraphTestUtils.waitUntilJobStatus(eg, JobStatus.RUNNING, timeout);

	ExecutionGraphTestUtils.waitForAllExecutionsPredicate(
		eg,
		isDeploying,
		timeout);

	//checking execution vertex properties
	validateConstraints(eg);

	ExecutionGraphTestUtils.finishAllVertices(eg);

	assertThat(eg.getState(), is(FINISHED));
}
 

@Test
public void testConstraintsAfterRestart() throws Exception {

	final long timeout = 5000L;

	//setting up
	testingSlotProvider.addTaskManager(NUM_TASKS);

	JobVertex groupVertex = ExecutionGraphTestUtils.createNoOpVertex(NUM_TASKS);
	JobVertex groupVertex2 = ExecutionGraphTestUtils.createNoOpVertex(NUM_TASKS);

	SlotSharingGroup sharingGroup = new SlotSharingGroup();
	groupVertex.setSlotSharingGroup(sharingGroup);
	groupVertex2.setSlotSharingGroup(sharingGroup);
	groupVertex.setStrictlyCoLocatedWith(groupVertex2);

	//initiate and schedule job
	final ExecutionGraph eg = new ExecutionGraphTestUtils.TestingExecutionGraphBuilder(groupVertex, groupVertex2)
		.setSlotProvider(testingSlotProvider)
		.setRestartStrategy(
			new TestRestartStrategy(
				1,
				false))
		.allowQueuedScheduling()
		.build();

	// enable the queued scheduling for the slot pool
	eg.start(ComponentMainThreadExecutorServiceAdapter.forMainThread());

	assertEquals(JobStatus.CREATED, eg.getState());

	eg.scheduleForExecution();

	Predicate<AccessExecution> isDeploying = ExecutionGraphTestUtils.isInExecutionState(ExecutionState.DEPLOYING);
	ExecutionGraphTestUtils.waitForAllExecutionsPredicate(
		eg,
		isDeploying,
		timeout);

	assertEquals(JobStatus.RUNNING, eg.getState());

	//sanity checks
	validateConstraints(eg);

	eg.getAllExecutionVertices().iterator().next().fail(new FlinkException("Test exception"));

	assertEquals(JobStatus.FAILING, eg.getState());

	for (ExecutionVertex vertex : eg.getAllExecutionVertices()) {
		vertex.getCurrentExecutionAttempt().completeCancelling();
	}

	// wait until we have restarted
	ExecutionGraphTestUtils.waitUntilJobStatus(eg, JobStatus.RUNNING, timeout);

	ExecutionGraphTestUtils.waitForAllExecutionsPredicate(
		eg,
		isDeploying,
		timeout);

	//checking execution vertex properties
	validateConstraints(eg);

	ExecutionGraphTestUtils.finishAllVertices(eg);

	assertThat(eg.getState(), is(FINISHED));
}
 

private void setSlotSharingAndCoLocation() {
	final HashMap<String, SlotSharingGroup> slotSharingGroups = new HashMap<>();
	final HashMap<String, Tuple2<SlotSharingGroup, CoLocationGroup>> coLocationGroups = new HashMap<>();

	for (Entry<Integer, JobVertex> entry : jobVertices.entrySet()) {

		final StreamNode node = streamGraph.getStreamNode(entry.getKey());
		final JobVertex vertex = entry.getValue();

		// configure slot sharing group
		final String slotSharingGroupKey = node.getSlotSharingGroup();
		final SlotSharingGroup sharingGroup;

		if (slotSharingGroupKey != null) {
			sharingGroup = slotSharingGroups.computeIfAbsent(
					slotSharingGroupKey, (k) -> new SlotSharingGroup());
			vertex.setSlotSharingGroup(sharingGroup);
		} else {
			sharingGroup = null;
		}

		// configure co-location constraint
		final String coLocationGroupKey = node.getCoLocationGroup();
		if (coLocationGroupKey != null) {
			if (sharingGroup == null) {
				throw new IllegalStateException("Cannot use a co-location constraint without a slot sharing group");
			}

			Tuple2<SlotSharingGroup, CoLocationGroup> constraint = coLocationGroups.computeIfAbsent(
					coLocationGroupKey, (k) -> new Tuple2<>(sharingGroup, new CoLocationGroup()));

			if (constraint.f0 != sharingGroup) {
				throw new IllegalStateException("Cannot co-locate operators from different slot sharing groups");
			}

			vertex.updateCoLocationGroup(constraint.f1);
		}
	}

	for (Tuple2<StreamNode, StreamNode> pair : streamGraph.getIterationSourceSinkPairs()) {

		CoLocationGroup ccg = new CoLocationGroup();

		JobVertex source = jobVertices.get(pair.f0.getId());
		JobVertex sink = jobVertices.get(pair.f1.getId());

		ccg.addVertex(source);
		ccg.addVertex(sink);
		source.updateCoLocationGroup(ccg);
		sink.updateCoLocationGroup(ccg);
	}

}
 

@Test
public void testRestartWithSlotSharingAndNotEnoughResources() throws Exception {
	// this test is inconclusive if not used with a proper multi-threaded executor
	assertTrue("test assumptions violated", ((ThreadPoolExecutor) executor).getCorePoolSize() > 1);

	final int numRestarts = 10;
	final int parallelism = 20;

	try (SlotPool slotPool = createSlotPoolImpl()) {
		final Scheduler scheduler = createSchedulerWithSlots(
			parallelism - 1, slotPool, new LocalTaskManagerLocation());

		final SlotSharingGroup sharingGroup = new SlotSharingGroup();

		final JobVertex source = new JobVertex("source");
		source.setInvokableClass(NoOpInvokable.class);
		source.setParallelism(parallelism);
		source.setSlotSharingGroup(sharingGroup);

		final JobVertex sink = new JobVertex("sink");
		sink.setInvokableClass(NoOpInvokable.class);
		sink.setParallelism(parallelism);
		sink.setSlotSharingGroup(sharingGroup);
		sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED_BOUNDED);

		TestRestartStrategy restartStrategy =
			new TestRestartStrategy(numRestarts, false);

		final ExecutionGraph eg = new ExecutionGraphTestUtils.TestingExecutionGraphBuilder(TEST_JOB_ID, source, sink)
			.setSlotProvider(scheduler)
			.setRestartStrategy(restartStrategy)
			.setIoExecutor(executor)
			.setFutureExecutor(executor)
			.setScheduleMode(ScheduleMode.EAGER)
			.build();

		eg.start(mainThreadExecutor);
		eg.scheduleForExecution();

		// wait until no more changes happen
		while (eg.getNumberOfFullRestarts() < numRestarts) {
			Thread.sleep(1);
		}

		assertEquals(JobStatus.FAILED, eg.getState());

		final Throwable t = eg.getFailureCause();
		if (!(t instanceof NoResourceAvailableException)) {
			ExceptionUtils.rethrowException(t, t.getMessage());
		}
	}
}
 

@Test
public void testRestartWithSlotSharingAndNotEnoughResources() throws Exception {
	final int numRestarts = 10;
	final int parallelism = 20;

	SlotPool slotPool = null;
	try {
		slotPool = new TestingSlotPoolImpl(TEST_JOB_ID);
		final Scheduler scheduler = createSchedulerWithSlots(
			parallelism - 1, slotPool, new LocalTaskManagerLocation());

		final SlotSharingGroup sharingGroup = new SlotSharingGroup();

		final JobVertex source = new JobVertex("source");
		source.setInvokableClass(NoOpInvokable.class);
		source.setParallelism(parallelism);
		source.setSlotSharingGroup(sharingGroup);

		final JobVertex sink = new JobVertex("sink");
		sink.setInvokableClass(NoOpInvokable.class);
		sink.setParallelism(parallelism);
		sink.setSlotSharingGroup(sharingGroup);
		sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED_BOUNDED);

		final JobGraph jobGraph = new JobGraph(TEST_JOB_ID, "Test Job", source, sink);
		jobGraph.setScheduleMode(ScheduleMode.EAGER);

		TestRestartStrategy restartStrategy = new TestRestartStrategy(numRestarts, false);

		final ExecutionGraph eg = TestingExecutionGraphBuilder
			.newBuilder()
			.setJobGraph(jobGraph)
			.setSlotProvider(scheduler)
			.setRestartStrategy(restartStrategy)
			.setAllocationTimeout(Time.milliseconds(1L))
			.build();

		eg.start(mainThreadExecutor);

		mainThreadExecutor.execute(ThrowingRunnable.unchecked(eg::scheduleForExecution));

		CommonTestUtils.waitUntilCondition(
			() -> CompletableFuture.supplyAsync(eg::getState, mainThreadExecutor).join() == JobStatus.FAILED,
			Deadline.fromNow(Duration.ofMillis(2000)));

		// the last suppressed restart is also counted
		assertEquals(numRestarts + 1, CompletableFuture.supplyAsync(eg::getNumberOfRestarts, mainThreadExecutor).join().longValue());

		final Throwable t = CompletableFuture.supplyAsync(eg::getFailureCause, mainThreadExecutor).join();
		if (!(t instanceof NoResourceAvailableException)) {
			ExceptionUtils.rethrowException(t, t.getMessage());
		}
	} finally {
		if (slotPool != null) {
			CompletableFuture.runAsync(slotPool::close, mainThreadExecutor).join();
		}
	}
}
 

@Test
public void testRestartWithEagerSchedulingAndSlotSharing() throws Exception {
	// this test is inconclusive if not used with a proper multi-threaded executor
	assertTrue("test assumptions violated", ((ThreadPoolExecutor) executor).getCorePoolSize() > 1);

	final int parallelism = 20;

	try (SlotPool slotPool = createSlotPoolImpl()) {
		final Scheduler scheduler = createSchedulerWithSlots(parallelism, slotPool, new LocalTaskManagerLocation());

		final SlotSharingGroup sharingGroup = new SlotSharingGroup();

		final JobVertex source = new JobVertex("source");
		source.setInvokableClass(NoOpInvokable.class);
		source.setParallelism(parallelism);
		source.setSlotSharingGroup(sharingGroup);

		final JobVertex sink = new JobVertex("sink");
		sink.setInvokableClass(NoOpInvokable.class);
		sink.setParallelism(parallelism);
		sink.setSlotSharingGroup(sharingGroup);
		sink.connectNewDataSetAsInput(source, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED_BOUNDED);

		TestRestartStrategy restartStrategy = TestRestartStrategy.directExecuting();

		final ExecutionGraph eg = new ExecutionGraphTestUtils.TestingExecutionGraphBuilder(TEST_JOB_ID, source, sink)
			.setSlotProvider(scheduler)
			.setIoExecutor(executor)
			.setFutureExecutor(executor)
			.setRestartStrategy(restartStrategy)
			.setScheduleMode(ScheduleMode.EAGER)
			.build();

		eg.start(mainThreadExecutor);

		eg.scheduleForExecution();

		switchToRunning(eg);

		// fail into 'RESTARTING'
		eg.getAllExecutionVertices().iterator().next().getCurrentExecutionAttempt().fail(
			new Exception("intended test failure"));

		assertEquals(JobStatus.FAILING, eg.getState());

		completeCancellingForAllVertices(eg);

		assertEquals(JobStatus.RUNNING, eg.getState());

		// clean termination
		switchToRunning(eg);
		finishAllVertices(eg);

		assertEquals(JobStatus.FINISHED, eg.getState());
	}
}
 

/**
 * Test setup:
 * - v1 is isolated, no slot sharing.
 * - v2 and v3 (not connected) share slots.
 * - v4 and v5 (connected) share slots.
 */
@Test
public void testAssignSlotSharingGroup() {
	try {
		JobVertex v1 = new JobVertex("v1");
		JobVertex v2 = new JobVertex("v2");
		JobVertex v3 = new JobVertex("v3");
		JobVertex v4 = new JobVertex("v4");
		JobVertex v5 = new JobVertex("v5");

		v1.setParallelism(4);
		v2.setParallelism(5);
		v3.setParallelism(7);
		v4.setParallelism(1);
		v5.setParallelism(11);

		v1.setInvokableClass(AbstractInvokable.class);
		v2.setInvokableClass(AbstractInvokable.class);
		v3.setInvokableClass(AbstractInvokable.class);
		v4.setInvokableClass(AbstractInvokable.class);
		v5.setInvokableClass(AbstractInvokable.class);

		v2.connectNewDataSetAsInput(v1, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);
		v5.connectNewDataSetAsInput(v4, DistributionPattern.POINTWISE, ResultPartitionType.PIPELINED);

		SlotSharingGroup jg1 = new SlotSharingGroup();
		v2.setSlotSharingGroup(jg1);
		v3.setSlotSharingGroup(jg1);

		SlotSharingGroup jg2 = new SlotSharingGroup();
		v4.setSlotSharingGroup(jg2);
		v5.setSlotSharingGroup(jg2);

		List<JobVertex> vertices = new ArrayList<>(Arrays.asList(v1, v2, v3, v4, v5));

		ExecutionGraph eg = TestingExecutionGraphBuilder.newBuilder().build();
		eg.attachJobGraph(vertices);

		// verify that the vertices are all in the same slot sharing group
		SlotSharingGroup group1;
		SlotSharingGroup group2;

		// verify that v1 tasks have no slot sharing group
		assertNull(eg.getJobVertex(v1.getID()).getSlotSharingGroup());

		// v2 and v3 are shared
		group1 = eg.getJobVertex(v2.getID()).getSlotSharingGroup();
		assertNotNull(group1);
		assertEquals(group1, eg.getJobVertex(v3.getID()).getSlotSharingGroup());

		assertEquals(2, group1.getJobVertexIds().size());
		assertTrue(group1.getJobVertexIds().contains(v2.getID()));
		assertTrue(group1.getJobVertexIds().contains(v3.getID()));

		// v4 and v5 are shared
		group2 = eg.getJobVertex(v4.getID()).getSlotSharingGroup();
		assertNotNull(group2);
		assertEquals(group2, eg.getJobVertex(v5.getID()).getSlotSharingGroup());

		assertEquals(2, group1.getJobVertexIds().size());
		assertTrue(group2.getJobVertexIds().contains(v4.getID()));
		assertTrue(group2.getJobVertexIds().contains(v5.getID()));
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests notifications of multiple receivers when a task produces both a pipelined and blocking
 * result.
 *
 * <pre>
 *                             +----------+
 *            +-- pipelined -> | Receiver |
 * +--------+ |                +----------+
 * | Sender |-|
 * +--------+ |                +----------+
 *            +-- blocking --> | Receiver |
 *                             +----------+
 * </pre>
 *
 * <p>The pipelined receiver gets deployed after the first buffer is available and the blocking
 * one after all subtasks are finished.
 */
@Test
public void testMixedPipelinedAndBlockingResults() throws Exception {
	final JobVertex sender = new JobVertex("Sender");
	sender.setInvokableClass(BinaryRoundRobinSubtaskIndexSender.class);
	sender.getConfiguration().setInteger(BinaryRoundRobinSubtaskIndexSender.CONFIG_KEY, PARALLELISM);
	sender.setParallelism(PARALLELISM);

	final JobVertex pipelinedReceiver = new JobVertex("Pipelined Receiver");
	pipelinedReceiver.setInvokableClass(SlotCountExceedingParallelismTest.SubtaskIndexReceiver.class);
	pipelinedReceiver.getConfiguration().setInteger(CONFIG_KEY, PARALLELISM);
	pipelinedReceiver.setParallelism(PARALLELISM);

	pipelinedReceiver.connectNewDataSetAsInput(
			sender,
			DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);

	final JobVertex blockingReceiver = new JobVertex("Blocking Receiver");
	blockingReceiver.setInvokableClass(SlotCountExceedingParallelismTest.SubtaskIndexReceiver.class);
	blockingReceiver.getConfiguration().setInteger(CONFIG_KEY, PARALLELISM);
	blockingReceiver.setParallelism(PARALLELISM);

	blockingReceiver.connectNewDataSetAsInput(sender,
			DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.BLOCKING);

	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();

	sender.setSlotSharingGroup(slotSharingGroup);
	pipelinedReceiver.setSlotSharingGroup(slotSharingGroup);
	blockingReceiver.setSlotSharingGroup(slotSharingGroup);

	final JobGraph jobGraph = new JobGraph(
			"Mixed pipelined and blocking result",
			sender,
			pipelinedReceiver,
			blockingReceiver);

	MINI_CLUSTER_RESOURCE.getMiniCluster().executeJobBlocking(jobGraph);
}
 

@Test
public void testSchedulingAllAtOnce() throws Exception {
	final int parallelism = 11;

	final MiniClusterConfiguration cfg = new MiniClusterConfiguration.Builder()
		.setNumTaskManagers(1)
		.setNumSlotsPerTaskManager(parallelism)
		.setConfiguration(getDefaultConfiguration())
		.build();

	try (final MiniCluster miniCluster = new MiniCluster(cfg)) {
		miniCluster.start();

		final JobVertex sender = new JobVertex("Sender");
		sender.setInvokableClass(Sender.class);
		sender.setParallelism(parallelism);

		final JobVertex forwarder = new JobVertex("Forwarder");
		forwarder.setInvokableClass(Forwarder.class);
		forwarder.setParallelism(parallelism);

		final JobVertex receiver = new JobVertex("Receiver");
		receiver.setInvokableClass(AgnosticReceiver.class);
		receiver.setParallelism(parallelism);

		final SlotSharingGroup sharingGroup = new SlotSharingGroup(sender.getID(), receiver.getID());
		sender.setSlotSharingGroup(sharingGroup);
		forwarder.setSlotSharingGroup(sharingGroup);
		receiver.setSlotSharingGroup(sharingGroup);

		forwarder.connectNewDataSetAsInput(sender, DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);
		receiver.connectNewDataSetAsInput(forwarder, DistributionPattern.ALL_TO_ALL,
			ResultPartitionType.PIPELINED);

		final JobGraph jobGraph = new JobGraph("Forwarding Job", sender, forwarder, receiver);

		jobGraph.setScheduleMode(ScheduleMode.EAGER);

		miniCluster.executeJobBlocking(jobGraph);
	}
}