Java Code Examples for org.apache.flink.runtime.io.network.partition.consumer.SingleInputGate#releaseAllResources()

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

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);
		inputChannel.requestSubpartition(0);

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), backlog);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), backlog);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline for notifying credits,
 * and verifies the behaviour of credit notification by triggering channel's writability changed.
 */
@Test
public void testNotifyCreditAvailable() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel1 = createRemoteInputChannel(inputGate, client);
	final RemoteInputChannel inputChannel2 = createRemoteInputChannel(inputGate, client);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		inputChannel1.requestSubpartition(0);
		inputChannel2.requestSubpartition(0);

		// The two input channels should send partition requests
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel1.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel2.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// The buffer response will take one available buffer from input channel, and it will trigger
		// requesting (backlog + numExclusiveBuffers - numAvailableBuffers) floating buffers
		final BufferResponse bufferResponse1 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel1.getInputChannelId(), 1);
		final BufferResponse bufferResponse2 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel2.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse2);

		assertEquals(2, inputChannel1.getUnannouncedCredit());
		assertEquals(2, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The two input channels should notify credits availability via the writable channel
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel1.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel2.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);
		assertNull(channel.readOutbound());

		ByteBuf channelBlockingBuffer = blockChannel(channel);

		// Trigger notify credits availability via buffer response on the condition of an un-writable channel
		final BufferResponse bufferResponse3 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 1, inputChannel1.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse3);

		assertEquals(1, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The input channel will not notify credits via un-writable channel
		assertFalse(channel.isWritable());
		assertNull(channel.readOutbound());

		// Flush the buffer to make the channel writable again
		channel.flush();
		assertSame(channelBlockingBuffer, channel.readOutbound());

		// The input channel should notify credits via channel's writability changed event
		assertTrue(channel.isWritable());
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(1, ((AddCredit) readFromOutbound).credit);
		assertEquals(0, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		// no more messages
		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline, but {@link AddCredit}
 * message is not sent actually when this input channel is released.
 */
@Test
public void testNotifyCreditAvailableAfterReleased() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		inputChannel.requestSubpartition(0);

		// This should send the partition request
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// Trigger request floating buffers via buffer response to notify credits available
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(2, inputChannel.getUnannouncedCredit());

		// Release the input channel
		inputGate.releaseAllResources();

		// it should send a close request after releasing the input channel,
		// but will not notify credits for a released input channel.
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(CloseRequest.class));

		channel.runPendingTasks();

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link NetworkEnvironment#registerTask(Task)} sets up (un)bounded buffer pool
 * instances for various types of input and output channels.
 */
@Test
public void testRegisterTaskUsesBoundedBuffers() throws Exception {
	final NetworkEnvironment network = new NetworkEnvironment(
		numBuffers, memorySegmentSize, 0, 0, 2, 8, enableCreditBasedFlowControl);

	// result partitions
	ResultPartition rp1 = createResultPartition(ResultPartitionType.PIPELINED, 2);
	ResultPartition rp2 = createResultPartition(ResultPartitionType.BLOCKING, 2);
	ResultPartition rp3 = createResultPartition(ResultPartitionType.PIPELINED_BOUNDED, 2);
	ResultPartition rp4 = createResultPartition(ResultPartitionType.PIPELINED_BOUNDED, 8);
	final ResultPartition[] resultPartitions = new ResultPartition[] {rp1, rp2, rp3, rp4};

	// input gates
	SingleInputGate ig1 = createSingleInputGate(ResultPartitionType.PIPELINED, 2);
	SingleInputGate ig2 = createSingleInputGate(ResultPartitionType.BLOCKING, 2);
	SingleInputGate ig3 = createSingleInputGate(ResultPartitionType.PIPELINED_BOUNDED, 2);
	SingleInputGate ig4 = createSingleInputGate(ResultPartitionType.PIPELINED_BOUNDED, 8);
	final SingleInputGate[] inputGates = new SingleInputGate[] {ig1, ig2, ig3, ig4};

	// overall task to register
	Task task = mock(Task.class);
	when(task.getProducedPartitions()).thenReturn(resultPartitions);
	when(task.getAllInputGates()).thenReturn(inputGates);

	network.registerTask(task);

	// verify buffer pools for the result partitions
	assertEquals(rp1.getNumberOfSubpartitions(), rp1.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(rp2.getNumberOfSubpartitions(), rp2.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(rp3.getNumberOfSubpartitions(), rp3.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(rp4.getNumberOfSubpartitions(), rp4.getBufferPool().getNumberOfRequiredMemorySegments());

	assertEquals(Integer.MAX_VALUE, rp1.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(Integer.MAX_VALUE, rp2.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(2 * 2 + 8, rp3.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(8 * 2 + 8, rp4.getBufferPool().getMaxNumberOfMemorySegments());

	// verify buffer pools for the input gates (NOTE: credit-based uses minimum required buffers
	// for exclusive buffers not managed by the buffer pool)
	assertEquals(enableCreditBasedFlowControl ? 0 : 2, ig1.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(enableCreditBasedFlowControl ? 0 : 2, ig2.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(enableCreditBasedFlowControl ? 0 : 2, ig3.getBufferPool().getNumberOfRequiredMemorySegments());
	assertEquals(enableCreditBasedFlowControl ? 0 : 8, ig4.getBufferPool().getNumberOfRequiredMemorySegments());

	assertEquals(Integer.MAX_VALUE, ig1.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(Integer.MAX_VALUE, ig2.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(enableCreditBasedFlowControl ? 8 : 2 * 2 + 8, ig3.getBufferPool().getMaxNumberOfMemorySegments());
	assertEquals(enableCreditBasedFlowControl ? 8 : 8 * 2 + 8, ig4.getBufferPool().getMaxNumberOfMemorySegments());

	int invokations = enableCreditBasedFlowControl ? 1 : 0;
	verify(ig1, times(invokations)).assignExclusiveSegments(network.getNetworkBufferPool(), 2);
	verify(ig2, times(invokations)).assignExclusiveSegments(network.getNetworkBufferPool(), 2);
	verify(ig3, times(invokations)).assignExclusiveSegments(network.getNetworkBufferPool(), 2);
	verify(ig4, times(invokations)).assignExclusiveSegments(network.getNetworkBufferPool(), 2);

	for (ResultPartition rp : resultPartitions) {
		rp.release();
	}
	for (SingleInputGate ig : inputGates) {
		ig.releaseAllResources();
	}
	network.shutdown();
}