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

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

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

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 testPartitionRequestSingleBackoff() throws Exception {
	// Setup
	SingleInputGate inputGate = createSingleInputGate(1);
	ResultPartitionID partitionId = new ResultPartitionID();
	TestVerifyPartitionRequestClient client = new TestVerifyPartitionRequestClient();
	ConnectionManager connectionManager = new TestVerifyConnectionManager(client);
	RemoteInputChannel ch = createRemoteInputChannel(inputGate, connectionManager, partitionId, 500, 500);

	// No delay for first request
	ch.requestSubpartition(0);
	client.verifyResult(partitionId, 0, 0);

	// Initial delay for second request
	ch.retriggerSubpartitionRequest(0);
	client.verifyResult(partitionId, 0, 500);

	// Exception after backoff is greater than the maximum backoff.
	try {
		ch.retriggerSubpartitionRequest(0);
		ch.getNextBuffer();
		fail("Did not throw expected exception.");
	}
	catch (Exception expected) {
	}
}
 

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

@Test
public void testPartitionRequestNoBackoff() throws Exception {
	// Setup
	SingleInputGate inputGate = createSingleInputGate(1);
	ResultPartitionID partitionId = new ResultPartitionID();
	TestVerifyPartitionRequestClient client = new TestVerifyPartitionRequestClient();
	ConnectionManager connectionManager = new TestVerifyConnectionManager(client);
	RemoteInputChannel ch = createRemoteInputChannel(inputGate, connectionManager, partitionId, 0, 0);

	// No delay for first request
	ch.requestSubpartition(0);
	client.verifyResult(partitionId, 0, 0);

	// Exception, because backoff is disabled.
	try {
		ch.retriggerSubpartitionRequest(0);
		ch.getNextBuffer();
		fail("Did not throw expected exception.");
	}
	catch (Exception expected) {
	}
}
 

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

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

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

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

@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 RemoteInputChannel ch = new RemoteInputChannel(
			mock(SingleInputGate.class),
			0,
			new ResultPartitionID(),
			mock(ConnectionID.class),
			connManager,
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

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

	ch.requestSubpartition(0);

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

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

/**
 * 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 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 RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		ConnectionManager connectionManager,
		ResultPartitionID partitionId,
		int initialBackoff,
		int maxBackoff) {
	return InputChannelBuilder.newBuilder()
		.setInitialBackoff(initialBackoff)
		.setMaxBackoff(maxBackoff)
		.setPartitionId(partitionId)
		.setConnectionManager(connectionManager)
		.buildRemoteChannel(inputGate);
}
 

@Test(expected = PartitionConnectionException.class)
public void testPartitionConnectionException() throws IOException {
	final ConnectionManager connManager = new TestingExceptionConnectionManager();
	final SingleInputGate gate = createSingleInputGate(1);
	final RemoteInputChannel ch = InputChannelTestUtils.createRemoteInputChannel(gate, 0, connManager);
	gate.setInputChannels(ch);

	gate.requestPartitions();

	ch.getNextBuffer();
}
 

public static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.buildRemoteChannel(inputGate);
}
 

public SingleInputGateBenchmarkFactory(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration networkConfig,
		ConnectionManager connectionManager,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		NetworkBufferPool networkBufferPool) {
	super(
		taskExecutorResourceId,
		networkConfig,
		connectionManager,
		partitionManager,
		taskEventPublisher,
		networkBufferPool);
}
 

public static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.buildRemoteChannel(inputGate);
}
 

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

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

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 static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.setMemorySegmentProvider(new UnpooledMemorySegmentProvider(32 * 1024))
		.buildRemoteAndSetToGate(inputGate);
}
 

public static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.buildRemoteAndSetToGate(inputGate);
}
 

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

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

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