org.apache.flink.runtime.io.network.partition.ResultPartitionManager Java Examples

public TestLocalInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		int initialBackoff,
		int maxBackoff,
		InputChannelMetrics metrics) {
	super(
		inputGate,
		channelIndex,
		partitionId,
		partitionManager,
		taskEventPublisher,
		initialBackoff,
		maxBackoff,
		metrics.getNumBytesInLocalCounter(),
		metrics.getNumBuffersInLocalCounter());
}
 

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(
			inputGate, partitionManager, new Tuple2<>(0, 0));

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

/**
 * Tests that an update channel does not trigger a partition request before the UDF has
 * requested any partitions. Otherwise, this can lead to races when registering a listener at
 * the gate (e.g. in UnionInputGate), which can result in missed buffer notifications at the
 * listener.
 */
@Test
public void testUpdateChannelBeforeRequest() throws Exception {
	SingleInputGate inputGate = createInputGate(1);

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	InputChannel unknown = new UnknownInputChannel(
		inputGate,
		0,
		new ResultPartitionID(),
		partitionManager,
		new TaskEventDispatcher(),
		new LocalConnectionManager(),
		0, 0, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	inputGate.setInputChannel(unknown.partitionId.getPartitionId(), unknown);

	// Update to a local channel and verify that no request is triggered
	inputGate.updateInputChannel(new InputChannelDeploymentDescriptor(
		unknown.partitionId,
		ResultPartitionLocation.createLocal()));

	verify(partitionManager, never()).createSubpartitionView(
		any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
}
 

public UnknownInputChannel(
		SingleInputGate gate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ConnectionManager connectionManager,
		int initialBackoff,
		int maxBackoff,
		TaskIOMetricGroup metrics) {

	super(gate, channelIndex, partitionId, initialBackoff, maxBackoff, null, null);

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
	this.connectionManager = checkNotNull(connectionManager);
	this.metrics = checkNotNull(metrics);
	this.initialBackoff = initialBackoff;
	this.maxBackoff = maxBackoff;
}
 

private LocalInputChannel createLocalInputChannel(
		SingleInputGate inputGate,
		ResultPartitionManager partitionManager,
		Tuple2<Integer, Integer> initialAndMaxRequestBackoff)
		throws IOException, InterruptedException {

	return new LocalInputChannel(
			inputGate,
			0,
			new ResultPartitionID(),
			partitionManager,
			mock(TaskEventDispatcher.class),
			initialAndMaxRequestBackoff._1(),
			initialAndMaxRequestBackoff._2(),
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
}
 

/**
 * Helper to create simple {@link ResultPartition} instance for use by a {@link Task} inside
 * {@link NetworkEnvironment#registerTask(Task)}.
 *
 * @param partitionType
 * 		the produced partition type
 * @param channels
 * 		the number of output channels
 *
 * @return instance with minimal data set and some mocks so that it is useful for {@link
 * NetworkEnvironment#registerTask(Task)}
 */
private static ResultPartition createResultPartition(
		final ResultPartitionType partitionType, final int channels) {
	return new ResultPartition(
		"TestTask-" + partitionType + ":" + channels,
		mock(TaskActions.class),
		new JobID(),
		new ResultPartitionID(),
		partitionType,
		channels,
		channels,
		mock(ResultPartitionManager.class),
		new NoOpResultPartitionConsumableNotifier(),
		mock(IOManager.class),
		false);
}
 

NettyShuffleEnvironment(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration config,
		NetworkBufferPool networkBufferPool,
		ConnectionManager connectionManager,
		ResultPartitionManager resultPartitionManager,
		FileChannelManager fileChannelManager,
		ResultPartitionFactory resultPartitionFactory,
		SingleInputGateFactory singleInputGateFactory) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.config = config;
	this.networkBufferPool = networkBufferPool;
	this.connectionManager = connectionManager;
	this.resultPartitionManager = resultPartitionManager;
	this.inputGatesById = new ConcurrentHashMap<>(10);
	this.fileChannelManager = fileChannelManager;
	this.resultPartitionFactory = resultPartitionFactory;
	this.singleInputGateFactory = singleInputGateFactory;
	this.isClosed = false;
}
 

public SingleInputGateFactory(
		@Nonnull ResourceID taskExecutorResourceId,
		@Nonnull NettyShuffleEnvironmentConfiguration networkConfig,
		@Nonnull ConnectionManager connectionManager,
		@Nonnull ResultPartitionManager partitionManager,
		@Nonnull TaskEventPublisher taskEventPublisher,
		@Nonnull NetworkBufferPool networkBufferPool) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.isCreditBased = networkConfig.isCreditBased();
	this.partitionRequestInitialBackoff = networkConfig.partitionRequestInitialBackoff();
	this.partitionRequestMaxBackoff = networkConfig.partitionRequestMaxBackoff();
	this.networkBuffersPerChannel = networkConfig.networkBuffersPerChannel();
	this.floatingNetworkBuffersPerGate = networkConfig.floatingNetworkBuffersPerGate();
	this.connectionManager = connectionManager;
	this.partitionManager = partitionManager;
	this.taskEventPublisher = taskEventPublisher;
	this.networkBufferPool = networkBufferPool;
}
 

public UnknownInputChannel(
		SingleInputGate gate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		ConnectionManager connectionManager,
		int initialBackoff,
		int maxBackoff,
		InputChannelMetrics metrics,
		@Nonnull MemorySegmentProvider memorySegmentProvider) {

	super(gate, channelIndex, partitionId, initialBackoff, maxBackoff, null, null);

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventPublisher = checkNotNull(taskEventPublisher);
	this.connectionManager = checkNotNull(connectionManager);
	this.metrics = checkNotNull(metrics);
	this.initialBackoff = initialBackoff;
	this.maxBackoff = maxBackoff;
	this.memorySegmentProvider = memorySegmentProvider;
}
 

/**
 * Tests that {@link PartitionRequestServerHandler} responds {@link ErrorResponse} with wrapped
 * {@link PartitionNotFoundException} after receiving invalid {@link PartitionRequest}.
 */
@Test
public void testResponsePartitionNotFoundException() {
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new PartitionRequestQueue(),
		true);
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	final ResultPartitionID partitionId = new ResultPartitionID();

	// Write the message of partition request to server
	channel.writeInbound(new PartitionRequest(partitionId, 0, new InputChannelID(), 2));
	channel.runPendingTasks();

	// Read the response message after handling partition request
	final Object msg = channel.readOutbound();
	assertThat(msg, instanceOf(ErrorResponse.class));

	final ErrorResponse err = (ErrorResponse) msg;
	assertThat(err.cause, instanceOf(PartitionNotFoundException.class));

	final ResultPartitionID actualPartitionId = ((PartitionNotFoundException) err.cause).getPartitionId();
	assertThat(partitionId, is(actualPartitionId));
}
 

/**
 * Tests that an update channel does not trigger a partition request before the UDF has
 * requested any partitions. Otherwise, this can lead to races when registering a listener at
 * the gate (e.g. in UnionInputGate), which can result in missed buffer notifications at the
 * listener.
 */
@Test
public void testUpdateChannelBeforeRequest() throws Exception {
	SingleInputGate inputGate = createInputGate(1);

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	InputChannel unknown = InputChannelBuilder.newBuilder()
		.setPartitionManager(partitionManager)
		.buildUnknownAndSetToGate(inputGate);

	// Update to a local channel and verify that no request is triggered
	ResultPartitionID resultPartitionID = unknown.getPartitionId();
	ResourceID location = ResourceID.generate();
	inputGate.updateInputChannel(location, createRemoteWithIdAndLocation(resultPartitionID.getPartitionId(), location));

	verify(partitionManager, never()).createSubpartitionView(
		any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
}
 

public NetworkEnvironment(
	int numBuffers,
	int memorySegmentSize,
	int partitionRequestInitialBackoff,
	int partitionRequestMaxBackoff,
	int networkBuffersPerChannel,
	int extraNetworkBuffersPerGate,
	boolean enableCreditBased) {
	this(
		new NetworkBufferPool(numBuffers, memorySegmentSize),
		new LocalConnectionManager(),
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new KvStateRegistry(),
		null,
		null,
		IOManager.IOMode.SYNC,
		partitionRequestInitialBackoff,
		partitionRequestMaxBackoff,
		networkBuffersPerChannel,
		extraNetworkBuffersPerGate,
		enableCreditBased);
}
 

NettyShuffleEnvironment(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration config,
		NetworkBufferPool networkBufferPool,
		ConnectionManager connectionManager,
		ResultPartitionManager resultPartitionManager,
		FileChannelManager fileChannelManager,
		ResultPartitionFactory resultPartitionFactory,
		SingleInputGateFactory singleInputGateFactory,
		Executor ioExecutor) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.config = config;
	this.networkBufferPool = networkBufferPool;
	this.connectionManager = connectionManager;
	this.resultPartitionManager = resultPartitionManager;
	this.inputGatesById = new ConcurrentHashMap<>(10);
	this.fileChannelManager = fileChannelManager;
	this.resultPartitionFactory = resultPartitionFactory;
	this.singleInputGateFactory = singleInputGateFactory;
	this.ioExecutor = ioExecutor;
	this.isClosed = false;
}
 

/**
 * Tests that if the {@link PartitionNotFoundException} is set onto one {@link InputChannel},
 * then it would be thrown directly via {@link SingleInputGate#getNextBufferOrEvent()}. So we
 * could confirm the {@link SingleInputGate} would not swallow or transform the original exception.
 */
@Test
public void testPartitionNotFoundExceptionWhileGetNextBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final LocalInputChannel localChannel = createLocalInputChannel(inputGate, new ResultPartitionManager());
	final ResultPartitionID partitionId = localChannel.getPartitionId();

	inputGate.setInputChannel(partitionId.getPartitionId(), localChannel);
	localChannel.setError(new PartitionNotFoundException(partitionId));
	try {
		inputGate.getNext();

		fail("Should throw a PartitionNotFoundException.");
	} catch (PartitionNotFoundException notFound) {
		assertThat(partitionId, is(notFound.getPartitionId()));
	}
}
 

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(inputGate, partitionManager);

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

/**
 * Tests that if the {@link PartitionNotFoundException} is set onto one {@link InputChannel},
 * then it would be thrown directly via {@link SingleInputGate#getNext()}. So we
 * could confirm the {@link SingleInputGate} would not swallow or transform the original exception.
 */
@Test
public void testPartitionNotFoundExceptionWhileGetNextBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final LocalInputChannel localChannel = createLocalInputChannel(inputGate, new ResultPartitionManager());
	final ResultPartitionID partitionId = localChannel.getPartitionId();

	inputGate.setInputChannels(localChannel);
	localChannel.setError(new PartitionNotFoundException(partitionId));
	try {
		inputGate.getNext();

		fail("Should throw a PartitionNotFoundException.");
	} catch (PartitionNotFoundException notFound) {
		assertThat(partitionId, is(notFound.getPartitionId()));
	}
}
 

@Test
public void testResumeConsumption() {
	final InputChannelID inputChannelID = new InputChannelID();
	final PartitionRequestQueue partitionRequestQueue = new PartitionRequestQueue();
	final TestViewReader testViewReader = new TestViewReader(inputChannelID, 2, partitionRequestQueue);
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		partitionRequestQueue);
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	partitionRequestQueue.notifyReaderCreated(testViewReader);

	// Write the message of resume consumption to server
	channel.writeInbound(new ResumeConsumption(inputChannelID));
	channel.runPendingTasks();

	assertTrue(testViewReader.consumptionResumed);
}
 

/**
 * Tests that {@link PartitionRequestServerHandler} responds {@link ErrorResponse} with wrapped
 * {@link PartitionNotFoundException} after receiving invalid {@link PartitionRequest}.
 */
@Test
public void testResponsePartitionNotFoundException() {
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new PartitionRequestQueue());
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	final ResultPartitionID partitionId = new ResultPartitionID();

	// Write the message of partition request to server
	channel.writeInbound(new PartitionRequest(partitionId, 0, new InputChannelID(), 2));
	channel.runPendingTasks();

	// Read the response message after handling partition request
	final Object msg = channel.readOutbound();
	assertThat(msg, instanceOf(ErrorResponse.class));

	final ErrorResponse err = (ErrorResponse) msg;
	assertThat(err.cause, instanceOf(PartitionNotFoundException.class));

	final ResultPartitionID actualPartitionId = ((PartitionNotFoundException) err.cause).getPartitionId();
	assertThat(partitionId, is(actualPartitionId));
}
 

@Test
public void testSlowIODoesNotBlockRelease() throws Exception {
	BlockerSync sync = new BlockerSync();
	ResultPartitionManager blockingResultPartitionManager = new ResultPartitionManager() {
		@Override
		public void releasePartition(ResultPartitionID partitionId, Throwable cause) {
			sync.blockNonInterruptible();
			super.releasePartition(partitionId, cause);
		}
	};

	NettyShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder()
		.setResultPartitionManager(blockingResultPartitionManager)
		.setIoExecutor(Executors.newFixedThreadPool(1))
		.build();

	shuffleEnvironment.releasePartitionsLocally(Collections.singleton(new ResultPartitionID()));
	sync.awaitBlocker();
	sync.releaseBlocker();
}
 

public UnknownInputChannel(
		SingleInputGate gate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		ConnectionManager connectionManager,
		int initialBackoff,
		int maxBackoff,
		InputChannelMetrics metrics) {

	super(gate, channelIndex, partitionId, initialBackoff, maxBackoff, null, null);

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventPublisher = checkNotNull(taskEventPublisher);
	this.connectionManager = checkNotNull(connectionManager);
	this.metrics = checkNotNull(metrics);
	this.initialBackoff = initialBackoff;
	this.maxBackoff = maxBackoff;
}
 

public SingleInputGateFactory(
		@Nonnull ResourceID taskExecutorResourceId,
		@Nonnull NettyShuffleEnvironmentConfiguration networkConfig,
		@Nonnull ConnectionManager connectionManager,
		@Nonnull ResultPartitionManager partitionManager,
		@Nonnull TaskEventPublisher taskEventPublisher,
		@Nonnull NetworkBufferPool networkBufferPool) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.partitionRequestInitialBackoff = networkConfig.partitionRequestInitialBackoff();
	this.partitionRequestMaxBackoff = networkConfig.partitionRequestMaxBackoff();
	this.networkBuffersPerChannel = networkConfig.networkBuffersPerChannel();
	this.floatingNetworkBuffersPerGate = networkConfig.floatingNetworkBuffersPerGate();
	this.blockingShuffleCompressionEnabled = networkConfig.isBlockingShuffleCompressionEnabled();
	this.compressionCodec = networkConfig.getCompressionCodec();
	this.networkBufferSize = networkConfig.networkBufferSize();
	this.connectionManager = connectionManager;
	this.partitionManager = partitionManager;
	this.taskEventPublisher = taskEventPublisher;
	this.networkBufferPool = networkBufferPool;
}
 

public LocalInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ResultPartitionManager partitionManager,
	TaskEventPublisher taskEventPublisher,
	int initialBackoff,
	int maxBackoff,
	Counter numBytesIn,
	Counter numBuffersIn) {

	super(inputGate, channelIndex, partitionId, initialBackoff, maxBackoff, numBytesIn, numBuffersIn);

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventPublisher = checkNotNull(taskEventPublisher);
}
 

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(inputGate, partitionManager);

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

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

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

	this.inputGate = new SingleInputGateBuilder()
		.setConsumedSubpartitionIndex(subpartitionIndex)
		.setNumberOfChannels(numberOfInputChannels)
		.setBufferPoolFactory(bufferPool)
		.build();

	// Setup input channels
	for (int i = 0; i < numberOfInputChannels; i++) {
		InputChannelBuilder.newBuilder()
			.setChannelIndex(i)
			.setPartitionManager(partitionManager)
			.setPartitionId(consumedPartitionIds[i])
			.setTaskEventPublisher(taskEventDispatcher)
			.buildLocalAndSetToGate(inputGate);
	}

	inputGate.setup();

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

public LocalInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ResultPartitionManager partitionManager,
	TaskEventDispatcher taskEventDispatcher,
	int initialBackoff,
	int maxBackoff,
	TaskIOMetricGroup metrics) {

	super(inputGate, channelIndex, partitionId, initialBackoff, maxBackoff, metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter());

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
}
 

@VisibleForTesting
static NettyShuffleEnvironment createNettyShuffleEnvironment(
		NettyShuffleEnvironmentConfiguration config,
		ResourceID taskExecutorResourceId,
		TaskEventPublisher taskEventPublisher,
		MetricGroup metricGroup,
		Executor ioExecutor) {
	return createNettyShuffleEnvironment(
		config,
		taskExecutorResourceId,
		taskEventPublisher,
		new ResultPartitionManager(),
		metricGroup,
		ioExecutor);
}
 

@Test
public void testExecutionFailsInNetworkRegistration() throws Exception {
	// mock a network manager that rejects registration
	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	final ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	final PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);

	final NetworkEnvironment network = mock(NetworkEnvironment.class);
	when(network.getResultPartitionManager()).thenReturn(partitionManager);
	when(network.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC);
	when(network.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);
	doThrow(new RuntimeException("buffers")).when(network).registerTask(any(Task.class));

	final QueuedNoOpTaskManagerActions taskManagerActions = new QueuedNoOpTaskManagerActions();
	final Task task = new TaskBuilder()
		.setTaskManagerActions(taskManagerActions)
		.setConsumableNotifier(consumableNotifier)
		.setPartitionProducerStateChecker(partitionProducerStateChecker)
		.setNetworkEnvironment(network)
		.build();

	// should fail
	task.run();

	// verify final state
	assertEquals(ExecutionState.FAILED, task.getExecutionState());
	assertTrue(task.isCanceledOrFailed());
	assertTrue(task.getFailureCause().getMessage().contains("buffers"));

	taskManagerActions.validateListenerMessage(
		ExecutionState.FAILED, task, new RuntimeException("buffers"));
}
 

/**
 * Tests that reading from a channel when after the partition has been
 * released are handled and don't lead to NPEs.
 */
@Test
public void testGetNextAfterPartitionReleased() throws Exception {
	ResultSubpartitionView reader = mock(ResultSubpartitionView.class);
	SingleInputGate gate = mock(SingleInputGate.class);
	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	when(partitionManager.createSubpartitionView(
		any(ResultPartitionID.class),
		anyInt(),
		any(BufferAvailabilityListener.class))).thenReturn(reader);

	LocalInputChannel channel = createLocalInputChannel(gate, partitionManager);

	channel.requestSubpartition(0);

	// Null buffer but not released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(false);

	assertFalse(channel.getNextBuffer().isPresent());

	// Null buffer and released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(true);

	try {
		channel.getNextBuffer();
		fail("Did not throw expected CancelTaskException");
	} catch (CancelTaskException ignored) {
	}

	channel.releaseAllResources();
	assertFalse(channel.getNextBuffer().isPresent());
}
 

/**
 * Tests that {@link LocalInputChannel#retriggerSubpartitionRequest(Timer, int)} would throw
 * {@link PartitionNotFoundException} which is set onto the input channel then.
 */
@Test
public void testChannelErrorWhileRetriggeringRequest() {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final LocalInputChannel localChannel = createLocalInputChannel(inputGate, new ResultPartitionManager());

	final Timer timer = new Timer(true) {
		@Override
		public void schedule(TimerTask task, long delay) {
			task.run();

			try {
				localChannel.checkError();

				fail("Should throw a PartitionNotFoundException.");
			} catch (PartitionNotFoundException notFound) {
				assertThat(localChannel.partitionId, Matchers.is(notFound.getPartitionId()));
			} catch (IOException ex) {
				fail("Should throw a PartitionNotFoundException.");
			}
		}
	};

	try {
		localChannel.retriggerSubpartitionRequest(timer, 0);
	} finally {
		timer.cancel();
	}
}
 

LocalRecoveredInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		int initialBackOff,
		int maxBackoff,
		InputChannelMetrics metrics) {
	super(inputGate, channelIndex, partitionId, initialBackOff, maxBackoff, metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter());

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventPublisher = checkNotNull(taskEventPublisher);
}