org.apache.flink.runtime.io.network.partition.consumer.SingleInputGate Java Examples

private CheckpointedInputGate createInputGate(
		int numberOfChannels,
		AbstractInvokable toNotify) throws IOException {
	final NettyShuffleEnvironment environment = new NettyShuffleEnvironmentBuilder().build();
	SingleInputGate gate = new SingleInputGateBuilder()
		.setNumberOfChannels(numberOfChannels)
		.setupBufferPoolFactory(environment)
		.build();
	gate.setInputChannels(
		IntStream.range(0, numberOfChannels)
		.mapToObj(channelIndex ->
			InputChannelBuilder.newBuilder()
				.setChannelIndex(channelIndex)
				.setupFromNettyShuffleEnvironment(environment)
				.setConnectionManager(new TestingConnectionManager())
				.buildRemoteChannel(gate))
		.toArray(RemoteInputChannel[]::new));
	sequenceNumbers = new int[numberOfChannels];

	gate.setup();
	gate.requestPartitions();

	return createCheckpointedInputGate(gate, toNotify);
}
 

public void registerTask(Task task) throws IOException {
	final ResultPartition[] producedPartitions = task.getProducedPartitions();

	synchronized (lock) {
		if (isShutdown) {
			throw new IllegalStateException("NetworkEnvironment is shut down");
		}

		for (final ResultPartition partition : producedPartitions) {
			setupPartition(partition);
		}

		// Setup the buffer pool for each buffer reader
		final SingleInputGate[] inputGates = task.getAllInputGates();
		for (SingleInputGate gate : inputGates) {
			setupInputGate(gate);
		}
	}
}
 

@Test
public void testOutOfOrderBarrier() throws Exception {
	SingleInputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	inputGate.setInputChannels(new TestInputChannel(inputGate, 0), new TestInputChannel(inputGate, 1));
	TestInvokable target = new TestInvokable();
	CheckpointBarrierAligner alignedHandler = new CheckpointBarrierAligner("test", target, inputGate);
	CheckpointBarrierUnaligner unalignedHandler = new CheckpointBarrierUnaligner(TestSubtaskCheckpointCoordinator.INSTANCE, "test", target, inputGate);
	AlternatingCheckpointBarrierHandler barrierHandler = new AlternatingCheckpointBarrierHandler(alignedHandler, unalignedHandler, target);

	long checkpointId = 10;
	long outOfOrderSavepointId = 5;
	long initialAlignedCheckpointId = alignedHandler.getLatestCheckpointId();

	barrierHandler.processBarrier(new CheckpointBarrier(checkpointId, 0, new CheckpointOptions(CHECKPOINT, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, 0));
	barrierHandler.processBarrier(new CheckpointBarrier(outOfOrderSavepointId, 0, new CheckpointOptions(SAVEPOINT, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, 1));

	assertEquals(checkpointId, barrierHandler.getLatestCheckpointId());
	assertInflightDataEquals(unalignedHandler, barrierHandler, checkpointId, inputGate.getNumberOfInputChannels());
	assertEquals(initialAlignedCheckpointId, alignedHandler.getLatestCheckpointId());
}
 

@Override
public Float getValue() {
	int usedBuffers = 0;
	int bufferPoolSize = 0;

	for (SingleInputGate inputGate : task.getAllInputGates()) {
		usedBuffers += inputGate.getBufferPool().bestEffortGetNumOfUsedBuffers();
		bufferPoolSize += inputGate.getBufferPool().getNumBuffers();
	}

	if (bufferPoolSize != 0) {
		return ((float) usedBuffers) / bufferPoolSize;
	} else {
		return 0.0f;
	}
}
 

@Before
public void setup() throws IOException, InterruptedException {
	CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	networkBufferPool = new NetworkBufferPool(
		NUMBER_OF_BUFFER_RESPONSES,
		BUFFER_SIZE,
		NUMBER_OF_BUFFER_RESPONSES);
	channel = new EmbeddedChannel(new NettyMessageClientDecoderDelegate(handler));

	inputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel inputChannel = createRemoteInputChannel(
		inputGate,
		new TestingPartitionRequestClient());
	inputGate.setInputChannels(inputChannel);
	inputGate.assignExclusiveSegments();
	inputChannel.requestSubpartition(0);
	handler.addInputChannel(inputChannel);
	inputChannelId = inputChannel.getInputChannelId();

	SingleInputGate releasedInputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel releasedInputChannel = new InputChannelBuilder()
		.buildRemoteChannel(inputGate);
	releasedInputGate.close();
	handler.addInputChannel(releasedInputChannel);
	releasedInputChannelId = releasedInputChannel.getInputChannelId();
}
 

/**
 * Verifies that {@link RemoteInputChannel#onError(Throwable)} is called when a
 * {@link BufferResponse} is received but no available buffer in input channel.
 */
@Test
public void testThrowExceptionForNoAvailableBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = spy(createRemoteInputChannel(inputGate));

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

	assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());

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

	verify(inputChannel, times(1)).onError(any(IllegalStateException.class));
}
 

private static void createRemoteInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartition resultPartition,
		ConnectionManager connManager) {
	RemoteInputChannel channel = new RemoteInputChannel(
		inputGate,
		channelIndex,
		resultPartition.getPartitionId(),
		mock(ConnectionID.class),
		connManager,
		0,
		0,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
	inputGate.setInputChannel(resultPartition.getPartitionId().getPartitionId(), channel);
}
 

@Test
public void testPreviousHandlerReset() throws Exception {
	SingleInputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	inputGate.setInputChannels(new TestInputChannel(inputGate, 0), new TestInputChannel(inputGate, 1));
	TestInvokable target = new TestInvokable();
	CheckpointBarrierAligner alignedHandler = new CheckpointBarrierAligner("test", target, inputGate);
	CheckpointBarrierUnaligner unalignedHandler = new CheckpointBarrierUnaligner(TestSubtaskCheckpointCoordinator.INSTANCE, "test", target, inputGate);
	AlternatingCheckpointBarrierHandler barrierHandler = new AlternatingCheckpointBarrierHandler(alignedHandler, unalignedHandler, target);

	for (int i = 0; i < 4; i++) {
		int channel = i % 2;
		CheckpointType type = channel == 0 ? CHECKPOINT : SAVEPOINT;
		barrierHandler.processBarrier(new CheckpointBarrier(i, 0, new CheckpointOptions(type, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, channel));
		assertEquals(type.isSavepoint(), alignedHandler.isCheckpointPending());
		assertNotEquals(alignedHandler.isCheckpointPending(), unalignedHandler.isCheckpointPending());

		if (!type.isSavepoint()) {
			assertFalse(barrierHandler.getAllBarriersReceivedFuture(i).isDone());
			assertInflightDataEquals(unalignedHandler, barrierHandler, i, inputGate.getNumberOfInputChannels());
		}
	}
}
 

@Override
public boolean updatePartitionInfo(
		ExecutionAttemptID consumerID,
		PartitionInfo partitionInfo) throws IOException, InterruptedException {
	IntermediateDataSetID intermediateResultPartitionID = partitionInfo.getIntermediateDataSetID();
	InputGateID id = new InputGateID(intermediateResultPartitionID, consumerID);
	SingleInputGate inputGate = inputGatesById.get(id);
	if (inputGate == null) {
		return false;
	}
	ShuffleDescriptor shuffleDescriptor = partitionInfo.getShuffleDescriptor();
	checkArgument(shuffleDescriptor instanceof NettyShuffleDescriptor,
		"Tried to update unknown channel with unknown ShuffleDescriptor %s.",
		shuffleDescriptor.getClass().getName());
	inputGate.updateInputChannel(taskExecutorResourceId, (NettyShuffleDescriptor) shuffleDescriptor);
	return true;
}
 

/**
 * Verifies that {@link RemoteInputChannel#onError(Throwable)} is called when a
 * {@link BufferResponse} is received but no available buffer in input channel.
 */
@Test
public void testThrowExceptionForNoAvailableBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = spy(InputChannelBuilder.newBuilder().buildRemoteAndSetToGate(inputGate));

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

	assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());

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

	verify(inputChannel, times(1)).onError(any(IllegalStateException.class));
}
 

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

@Override
public Float getValue() {
	int usedBuffers = 0;
	int totalBuffers = 0;

	for (SingleInputGate inputGate : inputGates) {
		usedBuffers += calculateUsedBuffers(inputGate);
		totalBuffers += calculateTotalBuffers(inputGate);
	}

	if (totalBuffers != 0) {
		return ((float) usedBuffers) / totalBuffers;
	} else {
		return 0.0f;
	}
}
 

@SuppressWarnings("unchecked")
public FairnessVerifyingInputGate(
		String owningTaskName,
		JobID jobId,
		IntermediateDataSetID consumedResultId,
		int consumedSubpartitionIndex,
		int numberOfInputChannels,
		TaskActions taskActions,
		TaskIOMetricGroup metrics,
		boolean isCreditBased) {

	super(owningTaskName, jobId, consumedResultId, ResultPartitionType.PIPELINED,
		consumedSubpartitionIndex,
			numberOfInputChannels, taskActions, metrics, isCreditBased);

	try {
		Field f = SingleInputGate.class.getDeclaredField("inputChannelsWithData");
		f.setAccessible(true);
		channelsWithData = (ArrayDeque<InputChannel>) f.get(this);
	}
	catch (Exception e) {
		throw new RuntimeException(e);
	}

	this.uniquenessChecker = new HashSet<>();
}
 

/**
 * Registers legacy network metric groups before shuffle service refactoring.
 *
 * <p>Registers legacy metric groups if shuffle service implementation is original default one.
 *
 * @deprecated should be removed in future
 */
@SuppressWarnings("DeprecatedIsStillUsed")
@Deprecated
public static void registerLegacyNetworkMetrics(
		boolean isDetailedMetrics,
		boolean isCreditBased,
		MetricGroup metricGroup,
		ResultPartitionWriter[] producedPartitions,
		InputGate[] inputGates) {
	checkNotNull(metricGroup);
	checkNotNull(producedPartitions);
	checkNotNull(inputGates);

	// add metrics for buffers
	final MetricGroup buffersGroup = metricGroup.addGroup(METRIC_GROUP_BUFFERS_DEPRECATED);

	// similar to MetricUtils.instantiateNetworkMetrics() but inside this IOMetricGroup (metricGroup)
	final MetricGroup networkGroup = metricGroup.addGroup(METRIC_GROUP_NETWORK_DEPRECATED);
	final MetricGroup outputGroup = networkGroup.addGroup(METRIC_GROUP_OUTPUT);
	final MetricGroup inputGroup = networkGroup.addGroup(METRIC_GROUP_INPUT);

	ResultPartition[] resultPartitions = Arrays.copyOf(producedPartitions, producedPartitions.length, ResultPartition[].class);
	registerOutputMetrics(isDetailedMetrics, outputGroup, buffersGroup, resultPartitions);

	SingleInputGate[] singleInputGates = Arrays.copyOf(inputGates, inputGates.length, SingleInputGate[].class);
	registerInputMetrics(isDetailedMetrics, isCreditBased, inputGroup, buffersGroup, singleInputGates);
}
 

private static RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartition resultPartition,
		ConnectionManager connManager) {
	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setPartitionId(resultPartition.getPartitionId())
		.setConnectionManager(connManager)
		.buildRemoteChannel(inputGate);
}
 

public static void registerInputMetrics(
		boolean isDetailedMetrics,
		boolean isCreditBased,
		MetricGroup inputGroup,
		SingleInputGate[] inputGates) {
	registerInputMetrics(
		isDetailedMetrics,
		isCreditBased,
		inputGroup,
		inputGroup.addGroup(METRIC_GROUP_BUFFERS),
		inputGates);
}
 

/**
 * 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, 2);
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder()
		.setMemorySegmentProvider(networkBufferPool)
		.buildRemoteAndSetToGate(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		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.close();

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

@Override
protected InputChannel createKnownInputChannel(
		SingleInputGate inputGate,
		int index,
		NettyShuffleDescriptor inputChannelDescriptor,
		SingleInputGateFactory.ChannelStatistics channelStatistics,
		InputChannelMetrics metrics) {
	ResultPartitionID partitionId = inputChannelDescriptor.getResultPartitionID();
	if (inputChannelDescriptor.isLocalTo(taskExecutorResourceId)) {
		return new TestLocalInputChannel(
			inputGate,
			index,
			partitionId,
			partitionManager,
			taskEventPublisher,
			partitionRequestInitialBackoff,
			partitionRequestMaxBackoff,
			metrics);
	} else {
		return new TestRemoteInputChannel(
			inputGate,
			index,
			partitionId,
			inputChannelDescriptor.getConnectionId(),
			connectionManager,
			partitionRequestInitialBackoff,
			partitionRequestMaxBackoff,
			metrics);
	}
}
 

@SuppressWarnings("unchecked")
public FairnessVerifyingInputGate(
		String owningTaskName,
		IntermediateDataSetID consumedResultId,
		int consumedSubpartitionIndex,
		int numberOfInputChannels) {

	super(
		owningTaskName,
		0,
		consumedResultId,
		ResultPartitionType.PIPELINED,
		consumedSubpartitionIndex,
		numberOfInputChannels,
		SingleInputGateBuilder.NO_OP_PRODUCER_CHECKER,
		STUB_BUFFER_POOL_FACTORY,
		null,
		new UnpooledMemorySegmentProvider(32 * 1024));

	try {
		Field f = SingleInputGate.class.getDeclaredField("inputChannelsWithData");
		f.setAccessible(true);
		channelsWithData = (ArrayDeque<InputChannel>) f.get(this);
	}
	catch (Exception e) {
		throw new RuntimeException(e);
	}

	this.uniquenessChecker = new HashSet<>();
}
 

@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 void registerInputMetrics(
		boolean isDetailedMetrics,
		MetricGroup inputGroup,
		SingleInputGate[] inputGates) {
	registerInputMetrics(
		isDetailedMetrics,
		inputGroup,
		inputGroup.addGroup(METRIC_GROUP_BUFFERS),
		inputGates);
}
 

@Override
public Integer getValue() {
	int totalBuffers = 0;

	for (SingleInputGate inputGate : task.getAllInputGates()) {
		totalBuffers += inputGate.getNumberOfQueuedBuffers();
	}

	return totalBuffers;
}
 

private StreamTaskNetworkInput<Long> createInput(CheckpointBarrierHandler handler, SingleInputGate inputGate) {
	return new StreamTaskNetworkInput<>(
		new CheckpointedInputGate(inputGate, handler),
		LongSerializer.INSTANCE,
		new StatusWatermarkValve(inputGate.getNumberOfInputChannels(), new NoOpDataOutput<>()),
		inputGate.getGateIndex(),
		createDeserializers(inputGate.getNumberOfInputChannels()));
}
 

public CreditBasedInputBuffersUsageGauge(
	FloatingBuffersUsageGauge floatingBuffersUsageGauge,
	ExclusiveBuffersUsageGauge exclusiveBuffersUsageGauge,
	SingleInputGate[] inputGates) {
	super(checkNotNull(inputGates));
	this.floatingBuffersUsageGauge = checkNotNull(floatingBuffersUsageGauge);
	this.exclusiveBuffersUsageGauge = checkNotNull(exclusiveBuffersUsageGauge);
}
 

public static LocalInputChannel createLocalInputChannel(
	SingleInputGate inputGate,
	ResultPartitionManager partitionManager,
	int initialBackoff,
	int maxBackoff) {

	return InputChannelBuilder.newBuilder()
		.setPartitionManager(partitionManager)
		.setInitialBackoff(initialBackoff)
		.setMaxBackoff(maxBackoff)
		.buildLocalChannel(inputGate);
}
 

public static LocalInputChannel createLocalInputChannel(
	SingleInputGate inputGate,
	ResultPartitionManager partitionManager,
	int initialBackoff,
	int maxBackoff) {

	return InputChannelBuilder.newBuilder()
		.setPartitionManager(partitionManager)
		.setInitialBackoff(initialBackoff)
		.setMaxBackoff(maxBackoff)
		.buildLocalAndSetToGate(inputGate);
}
 

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

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

public static void registerQueueLengthMetrics(MetricGroup parent, SingleInputGate[] gates) {
	for (int i = 0; i < gates.length; i++) {
		InputGateMetrics metrics = new InputGateMetrics(gates[i]);

		MetricGroup group = parent.addGroup(i);
		group.gauge("totalQueueLen", metrics.getTotalQueueLenGauge());
		group.gauge("minQueueLen", metrics.getMinQueueLenGauge());
		group.gauge("maxQueueLen", metrics.getMaxQueueLenGauge());
		group.gauge("avgQueueLen", metrics.getAvgQueueLenGauge());
	}
}
 

@Override
public int calculateTotalBuffers(SingleInputGate inputGate) {
	BufferPool bufferPool = inputGate.getBufferPool();
	if (bufferPool != null) {
		return inputGate.getBufferPool().getNumBuffers();
	}
	return 0;
}
 

@Override
public int calculateTotalBuffers(SingleInputGate inputGate) {
	BufferPool bufferPool = inputGate.getBufferPool();
	if (bufferPool != null) {
		return bufferPool.getNumBuffers();
	}
	return 0;
}