Java Code Examples for org.apache.flink.runtime.io.network.partition.ResultPartitionType#PIPELINED

private TaskDeploymentDescriptor createSender(
		NettyShuffleDescriptor shuffleDescriptor,
		Class<? extends AbstractInvokable> abstractInvokable) throws IOException {
	PartitionDescriptor partitionDescriptor = new PartitionDescriptor(
		new IntermediateDataSetID(),
		shuffleDescriptor.getResultPartitionID().getPartitionId(),
		ResultPartitionType.PIPELINED,
		1,
		0);
	ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = new ResultPartitionDeploymentDescriptor(
		partitionDescriptor,
		shuffleDescriptor,
		1,
		true);
	return createTestTaskDeploymentDescriptor(
		"Sender",
		shuffleDescriptor.getResultPartitionID().getProducerId(),
		abstractInvokable,
		1,
		Collections.singletonList(resultPartitionDeploymentDescriptor),
		Collections.emptyList());
}
 

@Test
public void testExecutionFailsInNetworkRegistrationForPartitions() throws Exception {
	final PartitionDescriptor partitionDescriptor = new PartitionDescriptor(
		new IntermediateDataSetID(),
		new IntermediateResultPartitionID(),
		ResultPartitionType.PIPELINED,
		1,
		1);
	final ShuffleDescriptor shuffleDescriptor = NettyShuffleDescriptorBuilder.newBuilder().buildLocal();
	final ResultPartitionDeploymentDescriptor dummyPartition = new ResultPartitionDeploymentDescriptor(
		partitionDescriptor,
		shuffleDescriptor,
		1,
		false);
	testExecutionFailsInNetworkRegistration(Collections.singleton(dummyPartition), Collections.emptyList());
}
 

private static ResultPartitionDeploymentDescriptor createResultPartitionDeploymentDescriptor(
		ResultPartitionID id,
		ResourceID location) throws ExecutionException, InterruptedException {
	ProducerDescriptor producerDescriptor = new ProducerDescriptor(
		location,
		id.getProducerId(),
		STUB_CONNECTION_ID.getAddress().getAddress(),
		STUB_CONNECTION_ID.getAddress().getPort());
	PartitionDescriptor partitionDescriptor = new PartitionDescriptor(
		new IntermediateDataSetID(),
		id.getPartitionId(),
		ResultPartitionType.PIPELINED,
		1,
		0);
	ShuffleDescriptor shuffleDescriptor =
		NettyShuffleMaster.INSTANCE.registerPartitionWithProducer(
			partitionDescriptor,
			producerDescriptor).get();
	return new ResultPartitionDeploymentDescriptor(
		partitionDescriptor,
		shuffleDescriptor,
		1,
		true);
}
 

/**
 * Tests that the lazy scheduling flag is correctly forwarded to the produced partition descriptors.
 */
@Test
public void testTddProducedPartitionsLazyScheduling() throws Exception {
	for (ScheduleMode scheduleMode: ScheduleMode.values()) {
		ExecutionJobVertex jobVertex = ExecutionGraphTestUtils.getExecutionJobVertex(
			new JobVertexID(),
			new DirectScheduledExecutorService(),
			scheduleMode);

		IntermediateResult result =
			new IntermediateResult(new IntermediateDataSetID(), jobVertex, 1, ResultPartitionType.PIPELINED);

		ExecutionAttemptID attemptID = new ExecutionAttemptID();
		ExecutionVertex vertex =
			new ExecutionVertex(jobVertex, 0, new IntermediateResult[]{result}, Time.minutes(1));
		TaskDeploymentDescriptorFactory tddFactory =
			TaskDeploymentDescriptorFactory.fromExecutionVertex(vertex, 1);

		ExecutionEdge mockEdge = createMockExecutionEdge(1);

		result.getPartitions()[0].addConsumerGroup();
		result.getPartitions()[0].addConsumer(mockEdge, 0);

		TaskManagerLocation location =
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1);

		TaskDeploymentDescriptor tdd = tddFactory.createDeploymentDescriptor(
			new AllocationID(),
			0,
			null,
			Execution.registerProducedPartitions(vertex, location, attemptID, scheduleMode.allowLazyDeployment()).get().values());

		Collection<ResultPartitionDeploymentDescriptor> producedPartitions = tdd.getProducedPartitions();

		assertEquals(1, producedPartitions.size());
		ResultPartitionDeploymentDescriptor desc = producedPartitions.iterator().next();
		assertEquals(scheduleMode.allowLazyDeployment(), desc.sendScheduleOrUpdateConsumersMessage());
	}
}
 

@Test
public void testExecutionFailsInNetworkRegistrationForGates() throws Exception {
	final ShuffleDescriptor dummyChannel = NettyShuffleDescriptorBuilder.newBuilder().buildRemote();
	final InputGateDeploymentDescriptor dummyGate = new InputGateDeploymentDescriptor(
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		new ShuffleDescriptor[] { dummyChannel });
	testExecutionFailsInNetworkRegistration(Collections.emptyList(), Collections.singletonList(dummyGate));
}
 

private static Map<InputGateID, SingleInputGate> createInputGateWithLocalChannels(
		NettyShuffleEnvironment network,
		int numberOfGates,
		@SuppressWarnings("SameParameterValue") int numberOfLocalChannels) {
	ShuffleDescriptor[] channelDescs = new NettyShuffleDescriptor[numberOfLocalChannels];
	for (int i = 0; i < numberOfLocalChannels; i++) {
		channelDescs[i] = createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), ResourceID.generate());
	}

	InputGateDeploymentDescriptor[] gateDescs = new InputGateDeploymentDescriptor[numberOfGates];
	IntermediateDataSetID[] ids = new IntermediateDataSetID[numberOfGates];
	for (int i = 0; i < numberOfGates; i++) {
		ids[i] = new IntermediateDataSetID();
		gateDescs[i] = new InputGateDeploymentDescriptor(
			ids[i],
			ResultPartitionType.PIPELINED,
			0,
			channelDescs);
	}

	ExecutionAttemptID consumerID = new ExecutionAttemptID();
	SingleInputGate[] gates = network.createInputGates(
		network.createShuffleIOOwnerContext("", consumerID, new UnregisteredMetricsGroup()),
		SingleInputGateBuilder.NO_OP_PRODUCER_CHECKER,
		Arrays.asList(gateDescs)).toArray(new SingleInputGate[] {});
	Map<InputGateID, SingleInputGate> inputGatesById = new HashMap<>();
	for (int i = 0; i < numberOfGates; i++) {
		inputGatesById.put(new InputGateID(ids[i], consumerID), gates[i]);
	}

	return inputGatesById;
}
 

/**
 * Tests that the lazy scheduling flag is correctly forwarded to the produced partition descriptors.
 */
@Test
public void testTddProducedPartitionsLazyScheduling() throws Exception {
	for (ScheduleMode scheduleMode: ScheduleMode.values()) {
		ExecutionJobVertex jobVertex = getExecutionVertex(
			new JobVertexID(),
			new DirectScheduledExecutorService(),
			scheduleMode);

		IntermediateResult result =
			new IntermediateResult(new IntermediateDataSetID(), jobVertex, 1, ResultPartitionType.PIPELINED);

		ExecutionAttemptID attemptID = new ExecutionAttemptID();
		ExecutionVertex vertex =
			new ExecutionVertex(jobVertex, 0, new IntermediateResult[]{result}, Time.minutes(1));
		TaskDeploymentDescriptorFactory tddFactory =
			TaskDeploymentDescriptorFactory.fromExecutionVertex(vertex, 1);

		ExecutionEdge mockEdge = createMockExecutionEdge(1);

		result.getPartitions()[0].addConsumerGroup();
		result.getPartitions()[0].addConsumer(mockEdge, 0);

		TaskManagerLocation location =
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1);

		TaskDeploymentDescriptor tdd = tddFactory.createDeploymentDescriptor(
			new AllocationID(),
			0,
			null,
			Execution.registerProducedPartitions(vertex, location, attemptID).get().values());

		Collection<ResultPartitionDeploymentDescriptor> producedPartitions = tdd.getProducedPartitions();

		assertEquals(1, producedPartitions.size());
		ResultPartitionDeploymentDescriptor desc = producedPartitions.iterator().next();
		assertEquals(scheduleMode.allowLazyDeployment(), desc.sendScheduleOrUpdateConsumersMessage());
	}
}
 

private TaskDeploymentDescriptor createReceiver(NettyShuffleDescriptor shuffleDescriptor) throws IOException {
	InputGateDeploymentDescriptor inputGateDeploymentDescriptor = new InputGateDeploymentDescriptor(
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		new ShuffleDescriptor[] {shuffleDescriptor});
	return createTestTaskDeploymentDescriptor(
		"Receiver",
		new ExecutionAttemptID(),
		TestingAbstractInvokables.Receiver.class,
		1,
		Collections.emptyList(),
		Collections.singletonList(inputGateDeploymentDescriptor));
}
 

@Test
public void testExecutionFailsInNetworkRegistrationForGates() throws Exception {
	final ShuffleDescriptor dummyChannel = NettyShuffleDescriptorBuilder.newBuilder().buildRemote();
	final InputGateDeploymentDescriptor dummyGate = new InputGateDeploymentDescriptor(
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		new ShuffleDescriptor[] { dummyChannel });
	testExecutionFailsInNetworkRegistration(Collections.emptyList(), Collections.singleton(dummyGate));
}
 

private SingleInputGate createSingleInputGate() {
	return new SingleInputGate(
		"InputGate",
		new JobID(),
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		1,
		mock(TaskActions.class),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
		true);
}
 

private static Map<InputGateID, SingleInputGate> createInputGateWithLocalChannels(
		NettyShuffleEnvironment network,
		int numberOfGates,
		@SuppressWarnings("SameParameterValue") int numberOfLocalChannels) {
	ShuffleDescriptor[] channelDescs = new NettyShuffleDescriptor[numberOfLocalChannels];
	for (int i = 0; i < numberOfLocalChannels; i++) {
		channelDescs[i] = createRemoteWithIdAndLocation(new IntermediateResultPartitionID(), ResourceID.generate());
	}

	InputGateDeploymentDescriptor[] gateDescs = new InputGateDeploymentDescriptor[numberOfGates];
	IntermediateDataSetID[] ids = new IntermediateDataSetID[numberOfGates];
	for (int i = 0; i < numberOfGates; i++) {
		ids[i] = new IntermediateDataSetID();
		gateDescs[i] = new InputGateDeploymentDescriptor(
			ids[i],
			ResultPartitionType.PIPELINED,
			0,
			channelDescs);
	}

	ExecutionAttemptID consumerID = new ExecutionAttemptID();
	SingleInputGate[] gates = network.createInputGates(
		network.createShuffleIOOwnerContext("", consumerID, new UnregisteredMetricsGroup()),
		SingleInputGateBuilder.NO_OP_PRODUCER_CHECKER,
		asList(gateDescs)).toArray(new SingleInputGate[] {});
	Map<InputGateID, SingleInputGate> inputGatesById = new HashMap<>();
	for (int i = 0; i < numberOfGates; i++) {
		inputGatesById.put(new InputGateID(ids[i], consumerID), gates[i]);
	}

	return inputGatesById;
}
 

/**
 * Tests that the lazy scheduling flag is correctly forwarded to the produced partition descriptors.
 */
@Test
public void testTddProducedPartitionsLazyScheduling() throws Exception {
	ExecutionJobVertex jobVertex = getExecutionVertex(new JobVertexID(), new DirectScheduledExecutorService());

	IntermediateResult result =
			new IntermediateResult(new IntermediateDataSetID(), jobVertex, 1, ResultPartitionType.PIPELINED);

	ExecutionVertex vertex =
			new ExecutionVertex(jobVertex, 0, new IntermediateResult[]{result}, Time.minutes(1));

	ExecutionEdge mockEdge = createMockExecutionEdge(1);

	result.getPartitions()[0].addConsumerGroup();
	result.getPartitions()[0].addConsumer(mockEdge, 0);

	SlotContext slotContext = mock(SlotContext.class);
	when(slotContext.getAllocationId()).thenReturn(new AllocationID());

	LogicalSlot slot = mock(LogicalSlot.class);
	when(slot.getAllocationId()).thenReturn(new AllocationID());

	for (ScheduleMode mode : ScheduleMode.values()) {
		vertex.getExecutionGraph().setScheduleMode(mode);

		TaskDeploymentDescriptor tdd = vertex.createDeploymentDescriptor(new ExecutionAttemptID(), slot, null, 1);

		Collection<ResultPartitionDeploymentDescriptor> producedPartitions = tdd.getProducedPartitions();

		assertEquals(1, producedPartitions.size());
		ResultPartitionDeploymentDescriptor desc = producedPartitions.iterator().next();
		assertEquals(mode.allowLazyDeployment(), desc.sendScheduleOrUpdateConsumersMessage());
	}
}
 

/**
 * Tests simple de/serialization.
 */
@Test
public void testSerialization() throws Exception {
	// Expected values
	IntermediateDataSetID resultId = new IntermediateDataSetID();
	IntermediateResultPartitionID partitionId = new IntermediateResultPartitionID();
	ResultPartitionType partitionType = ResultPartitionType.PIPELINED;
	int numberOfSubpartitions = 24;

	ResultPartitionDeploymentDescriptor orig =
			new ResultPartitionDeploymentDescriptor(
					resultId,
					partitionId,
					partitionType,
					numberOfSubpartitions,
					numberOfSubpartitions,
					true);

	ResultPartitionDeploymentDescriptor copy =
			CommonTestUtils.createCopySerializable(orig);

	assertEquals(resultId, copy.getResultId());
	assertEquals(partitionId, copy.getPartitionId());
	assertEquals(partitionType, copy.getPartitionType());
	assertEquals(numberOfSubpartitions, copy.getNumberOfSubpartitions());
	assertTrue(copy.sendScheduleOrUpdateConsumersMessage());
}
 

public TestLocalInputChannelConsumer(
		int subpartitionIndex,
		int numberOfInputChannels,
		int numberOfExpectedBuffersPerChannel,
		BufferPool bufferPool,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ResultPartitionID[] consumedPartitionIds) {

	checkArgument(numberOfInputChannels >= 1);
	checkArgument(numberOfExpectedBuffersPerChannel >= 1);

	this.inputGate = new SingleInputGate(
			"Test Name",
			new JobID(),
			new IntermediateDataSetID(),
			ResultPartitionType.PIPELINED,
			subpartitionIndex,
			numberOfInputChannels,
			mock(TaskActions.class),
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
			true);

	// Set buffer pool
	inputGate.setBufferPool(bufferPool);

	// Setup input channels
	for (int i = 0; i < numberOfInputChannels; i++) {
		inputGate.setInputChannel(
				new IntermediateResultPartitionID(),
				new LocalInputChannel(
						inputGate,
						i,
						consumedPartitionIds[i],
						partitionManager,
						taskEventDispatcher,
						UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup()));
	}

	this.numberOfInputChannels = numberOfInputChannels;
	this.numberOfExpectedBuffersPerChannel = numberOfExpectedBuffersPerChannel;
}
 

/**
 * Tests request back off configuration is correctly forwarded to the channels.
 */
@Test
public void testRequestBackoffConfiguration() throws Exception {
	ResultPartitionID[] partitionIds = new ResultPartitionID[] {
		new ResultPartitionID(),
		new ResultPartitionID(),
		new ResultPartitionID()
	};

	InputChannelDeploymentDescriptor[] channelDescs = new InputChannelDeploymentDescriptor[]{
		// Local
		new InputChannelDeploymentDescriptor(
			partitionIds[0],
			ResultPartitionLocation.createLocal()),
		// Remote
		new InputChannelDeploymentDescriptor(
			partitionIds[1],
			ResultPartitionLocation.createRemote(new ConnectionID(new InetSocketAddress("localhost", 5000), 0))),
		// Unknown
		new InputChannelDeploymentDescriptor(
			partitionIds[2],
			ResultPartitionLocation.createUnknown())};

	InputGateDeploymentDescriptor gateDesc =
		new InputGateDeploymentDescriptor(new IntermediateDataSetID(),
			ResultPartitionType.PIPELINED, 0, channelDescs);

	int initialBackoff = 137;
	int maxBackoff = 1001;

	final NetworkEnvironment netEnv = new NetworkEnvironment(
		100, 32, initialBackoff, maxBackoff, 2, 8, enableCreditBasedFlowControl);

	SingleInputGate gate = SingleInputGate.create(
		"TestTask",
		new JobID(),
		new ExecutionAttemptID(),
		gateDesc,
		netEnv,
		mock(TaskActions.class),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	try {
		assertEquals(gateDesc.getConsumedPartitionType(), gate.getConsumedPartitionType());

		Map<IntermediateResultPartitionID, InputChannel> channelMap = gate.getInputChannels();

		assertEquals(3, channelMap.size());
		InputChannel localChannel = channelMap.get(partitionIds[0].getPartitionId());
		assertEquals(LocalInputChannel.class, localChannel.getClass());

		InputChannel remoteChannel = channelMap.get(partitionIds[1].getPartitionId());
		assertEquals(RemoteInputChannel.class, remoteChannel.getClass());

		InputChannel unknownChannel = channelMap.get(partitionIds[2].getPartitionId());
		assertEquals(UnknownInputChannel.class, unknownChannel.getClass());

		InputChannel[] channels =
			new InputChannel[] {localChannel, remoteChannel, unknownChannel};
		for (InputChannel ch : channels) {
			assertEquals(0, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2 * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(maxBackoff, ch.getCurrentBackoff());

			assertFalse(ch.increaseBackoff());
		}
	} finally {
		gate.releaseAllResources();
		netEnv.shutdown();
	}
}
 

/**
 * Tests basic correctness of buffer-or-event interleaving and correct <code>null</code> return
 * value after receiving all end-of-partition events.
 *
 * <p>For buffer-or-event instances, it is important to verify that they have been set off to
 * the correct logical index.
 */
@Test(timeout = 120 * 1000)
public void testBasicGetNextLogic() throws Exception {
	// Setup
	final String testTaskName = "Test Task";
	final SingleInputGate ig1 = new SingleInputGate(
		testTaskName, new JobID(),
		new IntermediateDataSetID(), ResultPartitionType.PIPELINED,
		0, 3,
		mock(TaskActions.class),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
		true);
	final SingleInputGate ig2 = new SingleInputGate(
		testTaskName, new JobID(),
		new IntermediateDataSetID(), ResultPartitionType.PIPELINED,
		0, 5,
		mock(TaskActions.class),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
		true);

	final UnionInputGate union = new UnionInputGate(new SingleInputGate[]{ig1, ig2});

	assertEquals(ig1.getNumberOfInputChannels() + ig2.getNumberOfInputChannels(), union.getNumberOfInputChannels());

	final TestInputChannel[][] inputChannels = new TestInputChannel[][]{
			TestInputChannel.createInputChannels(ig1, 3),
			TestInputChannel.createInputChannels(ig2, 5)
	};

	inputChannels[0][0].readBuffer(); // 0 => 0
	inputChannels[0][0].readEndOfPartitionEvent(); // 0 => 0
	inputChannels[1][2].readBuffer(); // 2 => 5
	inputChannels[1][2].readEndOfPartitionEvent(); // 2 => 5
	inputChannels[1][0].readBuffer(); // 0 => 3
	inputChannels[1][1].readBuffer(); // 1 => 4
	inputChannels[0][1].readBuffer(); // 1 => 1
	inputChannels[1][3].readBuffer(); // 3 => 6
	inputChannels[0][1].readEndOfPartitionEvent(); // 1 => 1
	inputChannels[1][3].readEndOfPartitionEvent(); // 3 => 6
	inputChannels[0][2].readBuffer(); // 1 => 2
	inputChannels[0][2].readEndOfPartitionEvent(); // 1 => 2
	inputChannels[1][4].readBuffer(); // 4 => 7
	inputChannels[1][4].readEndOfPartitionEvent(); // 4 => 7
	inputChannels[1][1].readEndOfPartitionEvent(); // 0 => 3
	inputChannels[1][0].readEndOfPartitionEvent(); // 0 => 3

	ig1.notifyChannelNonEmpty(inputChannels[0][0]);
	ig1.notifyChannelNonEmpty(inputChannels[0][1]);
	ig1.notifyChannelNonEmpty(inputChannels[0][2]);

	ig2.notifyChannelNonEmpty(inputChannels[1][0]);
	ig2.notifyChannelNonEmpty(inputChannels[1][1]);
	ig2.notifyChannelNonEmpty(inputChannels[1][2]);
	ig2.notifyChannelNonEmpty(inputChannels[1][3]);
	ig2.notifyChannelNonEmpty(inputChannels[1][4]);

	verifyBufferOrEvent(union, true, 0, true); // gate 1, channel 0
	verifyBufferOrEvent(union, true, 3, true); // gate 2, channel 0
	verifyBufferOrEvent(union, true, 1, true); // gate 1, channel 1
	verifyBufferOrEvent(union, true, 4, true); // gate 2, channel 1
	verifyBufferOrEvent(union, true, 2, true); // gate 1, channel 2
	verifyBufferOrEvent(union, true, 5, true); // gate 2, channel 1
	verifyBufferOrEvent(union, false, 0, true); // gate 1, channel 0
	verifyBufferOrEvent(union, true, 6, true); // gate 2, channel 1
	verifyBufferOrEvent(union, false, 1, true); // gate 1, channel 1
	verifyBufferOrEvent(union, true, 7, true); // gate 2, channel 1
	verifyBufferOrEvent(union, false, 2, true); // gate 1, channel 2
	verifyBufferOrEvent(union, false, 3, true); // gate 2, channel 0
	verifyBufferOrEvent(union, false, 4, true); // gate 2, channel 1
	verifyBufferOrEvent(union, false, 5, true); // gate 2, channel 2
	verifyBufferOrEvent(union, false, 6, true); // gate 2, channel 3
	verifyBufferOrEvent(union, false, 7, false); // gate 2, channel 4

	// Return null when the input gate has received all end-of-partition events
	assertTrue(union.isFinished());
	assertFalse(union.getNextBufferOrEvent().isPresent());
}
 

public TestSingleInputGate(int numberOfInputChannels, boolean initialize) {
	checkArgument(numberOfInputChannels >= 1);

	SingleInputGate realGate = new SingleInputGate(
		"Test Task Name",
		new JobID(),
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		numberOfInputChannels,
		mock(TaskActions.class),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
		true);

	this.inputGate = spy(realGate);

	// Notify about late registrations (added for DataSinkTaskTest#testUnionDataSinkTask).
	// After merging registerInputOutput and invoke, we have to make sure that the test
	// notifications happen at the expected time. In real programs, this is guaranteed by
	// the instantiation and request partition life cycle.
	try {
		Field f = realGate.getClass().getDeclaredField("inputChannelsWithData");
		f.setAccessible(true);
		final ArrayDeque<InputChannel> notifications = (ArrayDeque<InputChannel>) f.get(realGate);

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) throws Throwable {
				invocation.callRealMethod();

				synchronized (notifications) {
					if (!notifications.isEmpty()) {
						InputGateListener listener = (InputGateListener) invocation.getArguments()[0];
						listener.notifyInputGateNonEmpty(inputGate);
					}
				}

				return null;
			}
		}).when(inputGate).registerListener(any(InputGateListener.class));
	} catch (Exception e) {
		throw new RuntimeException(e);
	}

	this.inputChannels = new TestInputChannel[numberOfInputChannels];

	if (initialize) {
		for (int i = 0; i < numberOfInputChannels; i++) {
			inputChannels[i] = new TestInputChannel(inputGate, i);
			inputGate.setInputChannel(new IntermediateResultPartitionID(), inputChannels[i]);
		}
	}
}
 

/**
 * Tests request back off configuration is correctly forwarded to the channels.
 */
@Test
public void testRequestBackoffConfiguration() throws Exception {
	IntermediateResultPartitionID[] partitionIds = new IntermediateResultPartitionID[] {
		new IntermediateResultPartitionID(),
		new IntermediateResultPartitionID(),
		new IntermediateResultPartitionID()
	};

	ResourceID localLocation = ResourceID.generate();
	ShuffleDescriptor[] channelDescs = new ShuffleDescriptor[]{
		// Local
		createRemoteWithIdAndLocation(partitionIds[0], localLocation),
		// Remote
		createRemoteWithIdAndLocation(partitionIds[1], ResourceID.generate()),
		// Unknown
		new UnknownShuffleDescriptor(new ResultPartitionID(partitionIds[2], new ExecutionAttemptID()))};

	InputGateDeploymentDescriptor gateDesc = new InputGateDeploymentDescriptor(
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		channelDescs);

	int initialBackoff = 137;
	int maxBackoff = 1001;

	final NettyShuffleEnvironment netEnv = new NettyShuffleEnvironmentBuilder()
		.setPartitionRequestInitialBackoff(initialBackoff)
		.setPartitionRequestMaxBackoff(maxBackoff)
		.setIsCreditBased(enableCreditBasedFlowControl)
		.build();

	SingleInputGate gate = new SingleInputGateFactory(
		localLocation,
		netEnv.getConfiguration(),
		netEnv.getConnectionManager(),
		netEnv.getResultPartitionManager(),
		new TaskEventDispatcher(),
		netEnv.getNetworkBufferPool())
		.create(
			"TestTask",
			gateDesc,
			SingleInputGateBuilder.NO_OP_PRODUCER_CHECKER,
			InputChannelTestUtils.newUnregisteredInputChannelMetrics());

	try {
		assertEquals(gateDesc.getConsumedPartitionType(), gate.getConsumedPartitionType());

		Map<IntermediateResultPartitionID, InputChannel> channelMap = gate.getInputChannels();

		assertEquals(3, channelMap.size());
		InputChannel localChannel = channelMap.get(partitionIds[0]);
		assertEquals(LocalInputChannel.class, localChannel.getClass());

		InputChannel remoteChannel = channelMap.get(partitionIds[1]);
		assertEquals(RemoteInputChannel.class, remoteChannel.getClass());

		InputChannel unknownChannel = channelMap.get(partitionIds[2]);
		assertEquals(UnknownInputChannel.class, unknownChannel.getClass());

		InputChannel[] channels =
			new InputChannel[] {localChannel, remoteChannel, unknownChannel};
		for (InputChannel ch : channels) {
			assertEquals(0, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2 * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(maxBackoff, ch.getCurrentBackoff());

			assertFalse(ch.increaseBackoff());
		}
	} finally {
		gate.close();
		netEnv.close();
	}
}
 

/**
 * Test that a failing schedule or update consumers call leads to the failing of the respective
 * task.
 *
 * <p>IMPORTANT: We have to make sure that the invokable's cancel method is called, because only
 * then the future is completed. We do this by not eagerly deploy consumer tasks and requiring
 * the invokable to fill one memory segment. The completed memory segment will trigger the
 * scheduling of the downstream operator since it is in pipeline mode. After we've filled the
 * memory segment, we'll block the invokable and wait for the task failure due to the failed
 * schedule or update consumers call.
 */
@Test(timeout = 10000L)
public void testFailingScheduleOrUpdateConsumersMessage() throws Exception {
	new JavaTestKit(system) {{
		final Configuration configuration = new Configuration();

		// set the memory segment to the smallest size possible, because we have to fill one
		// memory buffer to trigger the schedule or update consumers message to the downstream
		// operators
		configuration.setString(TaskManagerOptions.MEMORY_SEGMENT_SIZE, "4096");

		final JobID jid = new JobID();
		final JobVertexID vid = new JobVertexID();
		final ExecutionAttemptID eid = new ExecutionAttemptID();
		final SerializedValue<ExecutionConfig> executionConfig = new SerializedValue<>(new ExecutionConfig());

		final ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = new ResultPartitionDeploymentDescriptor(
			new IntermediateDataSetID(),
			new IntermediateResultPartitionID(),
			ResultPartitionType.PIPELINED,
			1,
			1,
			true);

		final TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(jid, "TestJob", vid, eid, executionConfig,
			"TestTask", 1, 0, 1, 0, new Configuration(), new Configuration(),
			TestInvokableRecordCancel.class.getName(),
			Collections.singletonList(resultPartitionDeploymentDescriptor),
			Collections.<InputGateDeploymentDescriptor>emptyList(),
			new ArrayList<>(), Collections.emptyList(), 0);

		ActorRef jmActorRef = system.actorOf(Props.create(FailingScheduleOrUpdateConsumersJobManager.class, LEADER_SESSION_ID), "jobmanager");
		ActorGateway jobManager = new AkkaActorGateway(jmActorRef, LEADER_SESSION_ID);

		highAvailabilityServices.setJobMasterLeaderRetriever(
			HighAvailabilityServices.DEFAULT_JOB_ID,
			new StandaloneLeaderRetrievalService(jobManager.path(), jobManager.leaderSessionID()));

		final ActorGateway taskManager = TestingUtils.createTaskManager(
			system,
			highAvailabilityServices,
			configuration,
			true,
			true);

		try {
			TestInvokableRecordCancel.resetGotCanceledFuture();

			Future<Object> result = taskManager.ask(new SubmitTask(tdd), timeout);

			Await.result(result, timeout);

			CompletableFuture<Boolean> cancelFuture = TestInvokableRecordCancel.gotCanceled();

			assertEquals(true, cancelFuture.get());
		} finally {
			TestingUtils.stopActor(taskManager);
			TestingUtils.stopActor(jobManager);
		}
	}};
}
 

private PartitionDescriptorBuilder() {
	this.partitionId = new IntermediateResultPartitionID();
	this.partitionType = ResultPartitionType.PIPELINED;
}