org.apache.flink.runtime.io.network.ConnectionID Java Examples

public RemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ConnectionID connectionId,
	ConnectionManager connectionManager,
	int initialBackOff,
	int maxBackoff,
	Counter numBytesIn,
	Counter numBuffersIn) {

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

	this.connectionId = checkNotNull(connectionId);
	this.connectionManager = checkNotNull(connectionManager);
	this.bufferManager = new BufferManager(inputGate.getMemorySegmentProvider(), this, 0);
}
 

/**
 * Creates and returns a remote input channel for the specific input gate with specific partition request client.
 *
 * @param inputGate The input gate owns the created input channel.
 * @param client The client is used to send partition request.
 * @param initialBackoff initial back off (in ms) for retriggering subpartition requests (must be <tt>&gt; 0</tt> to activate)
 * @param maxBackoff after which delay (in ms) to stop retriggering subpartition requests
 * @return The new created remote input channel.
 */
static RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		PartitionRequestClient client,
		int initialBackoff,
		int maxBackoff) throws Exception {
	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
		.thenReturn(client);

	ResultPartitionID partitionId = new ResultPartitionID();
	RemoteInputChannel inputChannel = new RemoteInputChannel(
		inputGate,
		0,
		partitionId,
		mock(ConnectionID.class),
		connectionManager,
		initialBackoff,
		maxBackoff,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	inputGate.setInputChannel(partitionId.getPartitionId(), inputChannel);
	return inputChannel;
}
 

public TestRemoteInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartitionID partitionId,
		ConnectionID connectionId,
		ConnectionManager connectionManager,
		int initialBackOff,
		int maxBackoff,
		InputChannelMetrics metrics) {
	super(
		inputGate,
		channelIndex,
		partitionId,
		connectionId,
		connectionManager,
		initialBackOff,
		maxBackoff,
		metrics.getNumBytesInRemoteCounter(),
		metrics.getNumBuffersInRemoteCounter());
}
 

@Test
public void testOnFailedPartitionRequest() throws Exception {
	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(mock(PartitionRequestClient.class));

	final ResultPartitionID partitionId = new ResultPartitionID();

	final SingleInputGate inputGate = mock(SingleInputGate.class);

	final RemoteInputChannel ch = new RemoteInputChannel(
			inputGate,
			0,
			partitionId,
			mock(ConnectionID.class),
			connectionManager,
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	ch.onFailedPartitionRequest();

	verify(inputGate).triggerPartitionStateCheck(eq(partitionId));
}
 

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {

	ConnectionManager connManager = mock(ConnectionManager.class);
	when(connManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(mock(PartitionRequestClient.class));

	final SingleInputGate gate = createSingleInputGate(1);
	final RemoteInputChannel ch = InputChannelTestUtils.createRemoteInputChannel(gate, 0, connManager);

	ch.onError(new ProducerFailedException(new RuntimeException("Expected test exception.")));

	ch.requestSubpartition(0);

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

private RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		PartitionRequestClient partitionRequestClient,
		Tuple2<Integer, Integer> initialAndMaxRequestBackoff)
		throws IOException, InterruptedException {

	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(partitionRequestClient);

	return new RemoteInputChannel(
		inputGate,
		0,
		new ResultPartitionID(),
		mock(ConnectionID.class),
		connectionManager,
		initialAndMaxRequestBackoff._1(),
		initialAndMaxRequestBackoff._2(),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
}
 

private RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		PartitionRequestClient partitionRequestClient,
		int initialBackoff,
		int maxBackoff,
		MemorySegmentProvider memorySegmentProvider)
		throws IOException, InterruptedException {

	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(partitionRequestClient);

	return InputChannelBuilder.newBuilder()
		.setConnectionManager(connectionManager)
		.setInitialBackoff(initialBackoff)
		.setMaxBackoff(maxBackoff)
		.setMemorySegmentProvider(memorySegmentProvider)
		.buildRemoteAndSetToGate(inputGate);
}
 

@Test
public void testOnFailedPartitionRequest() throws Exception {
	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(mock(PartitionRequestClient.class));

	final ResultPartitionID partitionId = new ResultPartitionID();

	final SingleInputGate inputGate = mock(SingleInputGate.class);

	final RemoteInputChannel ch = InputChannelBuilder.newBuilder()
		.setPartitionId(partitionId)
		.setConnectionManager(connectionManager)
		.buildRemoteAndSetToGate(inputGate);

	ch.onFailedPartitionRequest();

	verify(inputGate).triggerPartitionStateCheck(eq(partitionId));
}
 

/**
 * Creates a channel deployment descriptor by completing the partition location.
 *
 * @see InputChannelDeploymentDescriptor
 */
public InputChannelDeploymentDescriptor createInputChannelDeploymentDescriptor(Execution consumerExecution) {
	checkNotNull(consumerExecution, "consumerExecution");

	TaskManagerLocation consumerLocation = consumerExecution.getAssignedResourceLocation();
	checkNotNull(consumerLocation, "Consumer connection info null");

	final ResultPartitionLocation partitionLocation;

	if (consumerLocation.equals(partitionTaskManagerLocation)) {
		partitionLocation = ResultPartitionLocation.createLocal();
	}
	else {
		partitionLocation = ResultPartitionLocation.createRemote(
				new ConnectionID(partitionTaskManagerLocation, partitionConnectionIndex));
	}

	return new InputChannelDeploymentDescriptor(partitionID, partitionLocation);
}
 

public RemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ConnectionID connectionId,
	ConnectionManager connectionManager,
	int initialBackOff,
	int maxBackoff,
	InputChannelMetrics metrics,
	@Nonnull MemorySegmentProvider memorySegmentProvider) {

	super(inputGate, channelIndex, partitionId, initialBackOff, maxBackoff,
		metrics.getNumBytesInRemoteCounter(), metrics.getNumBuffersInRemoteCounter());

	this.connectionId = checkNotNull(connectionId);
	this.connectionManager = checkNotNull(connectionManager);
	this.memorySegmentProvider = memorySegmentProvider;
}
 

public RemoteInputChannel toRemoteInputChannel(ConnectionID producerAddress) {
	return new RemoteInputChannel(
		inputGate,
		getChannelIndex(),
		partitionId,
		checkNotNull(producerAddress),
		connectionManager,
		initialBackoff,
		maxBackoff,
		metrics.getNumBytesInRemoteCounter(),
		metrics.getNumBuffersInRemoteCounter());
}
 

private NettyPartitionRequestClient createPartitionRequestClient(
		Channel tcpChannel,
		NetworkClientHandler clientHandler) throws Exception {
	int port  = NetUtils.getAvailablePort();
	ConnectionID connectionID = new ConnectionID(new InetSocketAddress("localhost", port), 0);
	NettyConfig config = new NettyConfig(InetAddress.getLocalHost(), port, 1024, 1, new Configuration());
	NettyClient nettyClient = new NettyClient(config);
	PartitionRequestClientFactory partitionRequestClientFactory = new PartitionRequestClientFactory(nettyClient);

	return new NettyPartitionRequestClient(tcpChannel, clientHandler, connectionID, partitionRequestClientFactory);
}
 

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

public NettyShuffleDescriptor buildRemote() {
	ConnectionID connectionID = new ConnectionID(new InetSocketAddress(address, dataPort), connectionIndex);
	return new NettyShuffleDescriptor(
		producerLocation,
		new NetworkPartitionConnectionInfo(connectionID),
		id);
}
 

public void closeOpenChannelConnections(ConnectionID connectionId) {
	Object entry = clients.get(connectionId);

	if (entry instanceof ConnectingChannel) {
		ConnectingChannel channel = (ConnectingChannel) entry;

		if (channel.dispose()) {
			clients.remove(connectionId, channel);
		}
	}
}
 

RemoteRecoveredInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartitionID partitionId,
		ConnectionID connectionId,
		ConnectionManager connectionManager,
		int initialBackOff,
		int maxBackoff,
		InputChannelMetrics metrics) {
	super(inputGate, channelIndex, partitionId, initialBackOff, maxBackoff, metrics.getNumBytesInRemoteCounter(), metrics.getNumBuffersInRemoteCounter());

	this.connectionId = checkNotNull(connectionId);
	this.connectionManager = checkNotNull(connectionManager);
}
 

public RemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ConnectionID connectionId,
	ConnectionManager connectionManager,
	int initialBackOff,
	int maxBackoff,
	TaskIOMetricGroup metrics) {

	super(inputGate, channelIndex, partitionId, initialBackOff, maxBackoff, metrics.getNumBytesInRemoteCounter(), metrics.getNumBuffersInRemoteCounter());

	this.connectionId = checkNotNull(connectionId);
	this.connectionManager = checkNotNull(connectionManager);
}
 

PartitionRequestClient(
		Channel tcpChannel,
		NetworkClientHandler clientHandler,
		ConnectionID connectionId,
		PartitionRequestClientFactory clientFactory) {

	this.tcpChannel = checkNotNull(tcpChannel);
	this.clientHandler = checkNotNull(clientHandler);
	this.connectionId = checkNotNull(connectionId);
	this.clientFactory = checkNotNull(clientFactory);
}
 

public void closeOpenChannelConnections(ConnectionID connectionId) {
	Object entry = clients.get(connectionId);

	if (entry instanceof ConnectingChannel) {
		ConnectingChannel channel = (ConnectingChannel) entry;

		if (channel.dispose()) {
			clients.remove(connectionId, channel);
		}
	}
}
 

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

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

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		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.close();

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

public static ConnectionManager createDummyConnectionManager() throws Exception {
	final PartitionRequestClient mockClient = mock(PartitionRequestClient.class);

	final ConnectionManager connManager = mock(ConnectionManager.class);
	when(connManager.createPartitionRequestClient(any(ConnectionID.class))).thenReturn(mockClient);

	return connManager;
}
 

public static ConnectionManager createDummyConnectionManager() throws Exception {
	final PartitionRequestClient mockClient = mock(PartitionRequestClient.class);

	final ConnectionManager connManager = mock(ConnectionManager.class);
	when(connManager.createPartitionRequestClient(any(ConnectionID.class))).thenReturn(mockClient);

	return connManager;
}
 

static NetworkPartitionConnectionInfo fromProducerDescriptor(
		ProducerDescriptor producerDescriptor,
		int connectionIndex) {
	InetSocketAddress address =
		new InetSocketAddress(producerDescriptor.getAddress(), producerDescriptor.getDataPort());
	return new NetworkPartitionConnectionInfo(new ConnectionID(address, connectionIndex));
}
 

private void addRemoteInputChannel(
		NetworkEnvironment network,
		SingleInputGate inputGate,
		ConnectionID connectionId,
		ResultPartitionID partitionId,
		int channelIndex) {
	RemoteInputChannel remote =
		createUnknownInputChannel(network, inputGate, partitionId, channelIndex)
			.toRemoteInputChannel(connectionId);
	inputGate.setInputChannel(partitionId.getPartitionId(), remote);
}
 

@Test
public void testConsumptionWithRemoteChannels() throws Exception {
	final int numberOfChannels = 11;
	final int buffersPerChannel = 1000;

	final ConnectionManager connManager = createDummyConnectionManager();
	final Source[] sources = new Source[numberOfChannels];

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

	for (int i = 0; i < numberOfChannels; i++) {
		RemoteInputChannel channel = new RemoteInputChannel(
				gate, i, new ResultPartitionID(), mock(ConnectionID.class),
				connManager, 0, 0, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
		gate.setInputChannel(new IntermediateResultPartitionID(), channel);

		sources[i] = new RemoteChannelSource(channel);
	}

	ProducerThread producer = new ProducerThread(sources, numberOfChannels * buffersPerChannel, 4, 10);
	ConsumerThread consumer = new ConsumerThread(gate, numberOfChannels * buffersPerChannel);
	producer.start();
	consumer.start();

	// the 'sync()' call checks for exceptions and failed assertions
	producer.sync();
	consumer.sync();
}
 

static NetworkPartitionConnectionInfo fromProducerDescriptor(
		ProducerDescriptor producerDescriptor,
		int connectionIndex) {
	InetSocketAddress address =
		new InetSocketAddress(producerDescriptor.getAddress(), producerDescriptor.getDataPort());
	return new NetworkPartitionConnectionInfo(new ConnectionID(address, connectionIndex));
}
 

NettyPartitionRequestClient(
		Channel tcpChannel,
		NetworkClientHandler clientHandler,
		ConnectionID connectionId,
		PartitionRequestClientFactory clientFactory) {

	this.tcpChannel = checkNotNull(tcpChannel);
	this.clientHandler = checkNotNull(clientHandler);
	this.connectionId = checkNotNull(connectionId);
	this.clientFactory = checkNotNull(clientFactory);
}
 

public static ConnectionManager createDummyConnectionManager() throws Exception {
	final PartitionRequestClient mockClient = mock(PartitionRequestClient.class);

	final ConnectionManager connManager = mock(ConnectionManager.class);
	when(connManager.createPartitionRequestClient(any(ConnectionID.class))).thenReturn(mockClient);

	return connManager;
}
 

public void closeOpenChannelConnections(ConnectionID connectionId) {
	Object entry = clients.get(connectionId);

	if (entry instanceof ConnectingChannel) {
		ConnectingChannel channel = (ConnectingChannel) entry;

		if (channel.dispose()) {
			clients.remove(connectionId, channel);
		}
	}
}
 

public NettyShuffleDescriptor buildRemote() {
	ConnectionID connectionID = new ConnectionID(new InetSocketAddress(address, dataPort), connectionIndex);
	return new NettyShuffleDescriptor(
		producerLocation,
		new NetworkPartitionConnectionInfo(connectionID),
		id);
}