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

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

	for (ResultPartition resultPartition : task.getProducedPartitions()) {
		usedBuffers += resultPartition.getBufferPool().bestEffortGetNumOfUsedBuffers();
		bufferPoolSize += resultPartition.getBufferPool().getNumBuffers();
	}

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

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

@Override
public List<ResultPartition> createResultPartitionWriters(
		ShuffleIOOwnerContext ownerContext,
		List<ResultPartitionDeploymentDescriptor> resultPartitionDeploymentDescriptors) {
	synchronized (lock) {
		Preconditions.checkState(!isClosed, "The NettyShuffleEnvironment has already been shut down.");

		ResultPartition[] resultPartitions = new ResultPartition[resultPartitionDeploymentDescriptors.size()];
		for (int partitionIndex = 0; partitionIndex < resultPartitions.length; partitionIndex++) {
			resultPartitions[partitionIndex] = resultPartitionFactory.create(
				ownerContext.getOwnerName(),
				partitionIndex,
				resultPartitionDeploymentDescriptors.get(partitionIndex));
		}

		registerOutputMetrics(config.isNetworkDetailedMetrics(), ownerContext.getOutputGroup(), resultPartitions);
		return  Arrays.asList(resultPartitions);
	}
}
 

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

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

protected ResultPartitionWriter createResultPartition(
		JobID jobId,
		ResultPartitionID partitionId,
		NetworkEnvironment environment,
		int channels) throws Exception {

	ResultPartition resultPartition = new ResultPartition(
		"sender task",
		new NoOpTaskActions(),
		jobId,
		partitionId,
		ResultPartitionType.PIPELINED_BOUNDED,
		channels,
		1,
		environment.getResultPartitionManager(),
		new NoOpResultPartitionConsumableNotifier(),
		ioManager,
		false);

	environment.setupPartition(resultPartition);

	return resultPartition;
}
 

@Override
public Collection<ResultPartition> createResultPartitionWriters(
		ShuffleIOOwnerContext ownerContext,
		Collection<ResultPartitionDeploymentDescriptor> resultPartitionDeploymentDescriptors) {
	synchronized (lock) {
		Preconditions.checkState(!isClosed, "The NettyShuffleEnvironment has already been shut down.");

		ResultPartition[] resultPartitions = new ResultPartition[resultPartitionDeploymentDescriptors.size()];
		int counter = 0;
		for (ResultPartitionDeploymentDescriptor rpdd : resultPartitionDeploymentDescriptors) {
			resultPartitions[counter++] = resultPartitionFactory.create(ownerContext.getOwnerName(), rpdd);
		}

		registerOutputMetrics(config.isNetworkDetailedMetrics(), ownerContext.getOutputGroup(), resultPartitions);
		return  Arrays.asList(resultPartitions);
	}
}
 

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

	for (ResultPartition resultPartition : resultPartitions) {
		BufferPool bufferPool = resultPartition.getBufferPool();

		if (bufferPool != null) {
			usedBuffers += bufferPool.bestEffortGetNumOfUsedBuffers();
			bufferPoolSize += bufferPool.getNumBuffers();
		}
	}

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

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

	for (ResultPartition resultPartition : resultPartitions) {
		BufferPool bufferPool = resultPartition.getBufferPool();

		if (bufferPool != null) {
			usedBuffers += bufferPool.bestEffortGetNumOfUsedBuffers();
			bufferPoolSize += bufferPool.getNumBuffers();
		}
	}

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

/**
 * 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,
		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, inputGroup, buffersGroup, singleInputGates);
}
 

private static ResultPartition createFinishedPartitionWithFilledData(ResultPartitionManager partitionManager) throws Exception {
	final ResultPartition partition = new ResultPartitionBuilder()
		.setResultPartitionType(ResultPartitionType.BLOCKING)
		.setFileChannelManager(fileChannelManager)
		.setResultPartitionManager(partitionManager)
		.isReleasedOnConsumption(true)
		.build();

	partitionManager.registerResultPartition(partition);
	PartitionTestUtils.writeBuffers(partition, 1, BUFFER_SIZE);

	return partition;
}
 

private RemoteInputChannel buildRemoteChannel(
	int channelIndex,
	SingleInputGate inputGate,
	NettyShuffleEnvironment network,
	ResultPartition partition) {
	return new InputChannelBuilder()
		.setPartitionId(partition.getPartitionId())
		.setChannelIndex(channelIndex)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildRemoteAndSetToGate(inputGate);
}
 

private void buildLocalChannel(
	int channelIndex,
	SingleInputGate inputGate,
	NettyShuffleEnvironment network,
	ResultPartition partition) {
	new InputChannelBuilder()
		.setPartitionId(partition.getPartitionId())
		.setChannelIndex(channelIndex)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildLocalAndSetToGate(inputGate);
}
 

public TestPartitionProducer(
		ResultPartition partition,
		boolean isSlowProducer,
		TestProducerSource source) {

	this.partition = checkNotNull(partition);
	this.isSlowProducer = isSlowProducer;
	this.random = isSlowProducer ? new Random() : null;
	this.source = checkNotNull(source);
}
 

public static void registerQueueLengthMetrics(MetricGroup parent, ResultPartition[] partitions) {
	for (int i = 0; i < partitions.length; i++) {
		ResultPartitionMetrics metrics = new ResultPartitionMetrics(partitions[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 Integer getValue() {
	int totalBuffers = 0;

	for (ResultPartition producedPartition : resultPartitions) {
		totalBuffers += producedPartition.getNumberOfQueuedBuffers();
	}

	return totalBuffers;
}
 

public static void registerOutputMetrics(
		boolean isDetailedMetrics,
		MetricGroup outputGroup,
		ResultPartition[] resultPartitions) {
	registerOutputMetrics(
		isDetailedMetrics,
		outputGroup,
		outputGroup.addGroup(METRIC_GROUP_BUFFERS),
		resultPartitions);
}
 

private static void registerOutputMetrics(
		boolean isDetailedMetrics,
		MetricGroup outputGroup,
		MetricGroup buffersGroup,
		ResultPartition[] resultPartitions) {
	if (isDetailedMetrics) {
		ResultPartitionMetrics.registerQueueLengthMetrics(outputGroup, resultPartitions);
	}
	buffersGroup.gauge(METRIC_OUTPUT_QUEUE_LENGTH, new OutputBuffersGauge(resultPartitions));
	buffersGroup.gauge(METRIC_OUTPUT_POOL_USAGE, new OutputBufferPoolUsageGauge(resultPartitions));
}
 

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

private void testCancelPartitionRequest(boolean isAvailableView) throws Exception {
	// setup
	final ResultPartitionManager partitionManager = new ResultPartitionManager();
	final ResultPartition partition = createFinishedPartitionWithFilledData(partitionManager);
	final InputChannelID receiverId = new InputChannelID();
	final PartitionRequestQueue queue = new PartitionRequestQueue();
	final CreditBasedSequenceNumberingViewReader reader = new CreditBasedSequenceNumberingViewReader(receiverId, 0, queue);
	final EmbeddedChannel channel = new EmbeddedChannel(queue);

	reader.requestSubpartitionView(partitionManager, partition.getPartitionId(), 0);
	// add this reader into allReaders queue
	queue.notifyReaderCreated(reader);

	// block the channel so that we see an intermediate state in the test
	blockChannel(channel);

	// add credit to make this reader available for adding into availableReaders queue
	if (isAvailableView) {
		queue.addCreditOrResumeConsumption(receiverId, viewReader -> viewReader.addCredit(1));
		assertTrue(queue.getAvailableReaders().contains(reader));
	}

	// cancel this subpartition view
	queue.cancel(receiverId);
	channel.runPendingTasks();

	assertFalse(queue.getAvailableReaders().contains(reader));
	// the partition and its reader view should all be released
	assertTrue(reader.isReleased());
	assertTrue(partition.isReleased());
	for (ResultSubpartition subpartition : partition.getAllPartitions()) {
		assertTrue(subpartition.isReleased());
	}

	// cleanup
	channel.close();
}
 

private static ResultPartition createFinishedPartitionWithFilledData(ResultPartitionManager partitionManager) throws Exception {
	final ResultPartition partition = new ResultPartitionBuilder()
		.setResultPartitionType(ResultPartitionType.BLOCKING)
		.setFileChannelManager(fileChannelManager)
		.setResultPartitionManager(partitionManager)
		.isReleasedOnConsumption(true)
		.build();

	partitionManager.registerResultPartition(partition);
	PartitionTestUtils.writeBuffers(partition, 1, BUFFER_SIZE);

	return partition;
}
 

private RemoteInputChannel buildRemoteChannel(
	int channelIndex,
	SingleInputGate inputGate,
	NettyShuffleEnvironment network,
	ResultPartition partition) {
	return new InputChannelBuilder()
		.setPartitionId(partition.getPartitionId())
		.setChannelIndex(channelIndex)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildRemoteChannel(inputGate);
}
 

private LocalInputChannel buildLocalChannel(
	int channelIndex,
	SingleInputGate inputGate,
	NettyShuffleEnvironment network,
	ResultPartition partition) {
	return new InputChannelBuilder()
		.setPartitionId(partition.getPartitionId())
		.setChannelIndex(channelIndex)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildLocalChannel(inputGate);
}
 

private static ResultSubpartitionView createResultSubpartitionView(boolean addBuffer) throws IOException {
	int bufferSize = 4096;
	ResultPartition parent = PartitionTestUtils.createPartition(
		ResultPartitionType.PIPELINED,
		NoOpFileChannelManager.INSTANCE,
		true,
		bufferSize);
	ResultSubpartition subpartition = parent.getAllPartitions()[0];
	if (addBuffer) {
		subpartition.add(BufferBuilderTestUtils.createFilledFinishedBufferConsumer(bufferSize));
	}
	return subpartition.createReadView(() -> {});
}
 

public TestPartitionProducer(
		ResultPartition partition,
		boolean isSlowProducer,
		TestProducerSource source) {

	this.partition = checkNotNull(partition);
	this.isSlowProducer = isSlowProducer;
	this.random = isSlowProducer ? new Random() : null;
	this.source = checkNotNull(source);
}
 

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

	for (ResultPartition producedPartition : resultPartitions) {
		totalBuffers += producedPartition.getNumberOfQueuedBuffers();
	}

	return totalBuffers;
}
 

public TaskCanceler(
		Logger logger,
		AbstractInvokable invokable,
		Thread executer,
		String taskName,
		ResultPartition[] producedPartitions,
		SingleInputGate[] inputGates) {

	this.logger = logger;
	this.invokable = invokable;
	this.executer = executer;
	this.taskName = taskName;
	this.producedPartitions = producedPartitions;
	this.inputGates = inputGates;
}
 

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

	for (ResultPartition producedPartition : task.getProducedPartitions()) {
		totalBuffers += producedPartition.getNumberOfQueuedBuffers();
	}

	return totalBuffers;
}
 

@VisibleForTesting
public void setupPartition(ResultPartition partition) throws IOException {
	BufferPool bufferPool = null;

	try {
		int maxNumberOfMemorySegments = partition.getPartitionType().isBounded() ?
			partition.getNumberOfSubpartitions() * networkBuffersPerChannel +
				extraNetworkBuffersPerGate : Integer.MAX_VALUE;
		// If the partition type is back pressure-free, we register with the buffer pool for
		// callbacks to release memory.
		bufferPool = networkBufferPool.createBufferPool(partition.getNumberOfSubpartitions(),
			maxNumberOfMemorySegments,
			partition.getPartitionType().hasBackPressure() ? Optional.empty() : Optional.of(partition));

		partition.registerBufferPool(bufferPool);

		resultPartitionManager.registerResultPartition(partition);
	} catch (Throwable t) {
		if (bufferPool != null) {
			bufferPool.lazyDestroy();
		}

		if (t instanceof IOException) {
			throw (IOException) t;
		} else {
			throw new IOException(t.getMessage(), t);
		}
	}

	taskEventDispatcher.registerPartition(partition.getPartitionId());
}
 

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