org.apache.flink.runtime.io.network.partition.PartitionTrackerImpl Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of flink-runtime_2.11 Show documentation
There is a newer version: 1.13.6
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.io.network.partition;

import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleMaster;
import org.apache.flink.runtime.taskexecutor.partition.PartitionTable;
import org.apache.flink.util.Preconditions;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import static java.util.stream.Collectors.toList;

/**
 * Utility for tracking partitions and issuing release calls to task executors and shuffle masters.
 */
public class PartitionTrackerImpl implements PartitionTracker {

	private final JobID jobId;

	private final PartitionTable partitionTable = new PartitionTable<>();
	private final Map partitionInfos = new HashMap<>();

	private final ShuffleMaster shuffleMaster;

	private final PartitionTrackerFactory.TaskExecutorGatewayLookup taskExecutorGatewayLookup;

	public PartitionTrackerImpl(
		JobID jobId,
		ShuffleMaster shuffleMaster,
		PartitionTrackerFactory.TaskExecutorGatewayLookup taskExecutorGatewayLookup) {

		this.jobId = Preconditions.checkNotNull(jobId);
		this.shuffleMaster = Preconditions.checkNotNull(shuffleMaster);
		this.taskExecutorGatewayLookup = taskExecutorGatewayLookup;
	}

	@Override
	public void startTrackingPartition(ResourceID producingTaskExecutorId, ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor) {
		Preconditions.checkNotNull(producingTaskExecutorId);
		Preconditions.checkNotNull(resultPartitionDeploymentDescriptor);

		// only blocking partitions require explicit release call
		if (!resultPartitionDeploymentDescriptor.getPartitionType().isBlocking()) {
			return;
		}

		final ResultPartitionID resultPartitionId = resultPartitionDeploymentDescriptor.getShuffleDescriptor().getResultPartitionID();

		partitionInfos.put(resultPartitionId, new PartitionInfo(producingTaskExecutorId, resultPartitionDeploymentDescriptor));
		partitionTable.startTrackingPartitions(producingTaskExecutorId, Collections.singletonList(resultPartitionId));
	}

	@Override
	public void stopTrackingPartitionsFor(ResourceID producingTaskExecutorId) {
		Preconditions.checkNotNull(producingTaskExecutorId);

		// this is a bit icky since we make 2 calls to pT#stopTrackingPartitions
		final Collection resultPartitionIds = partitionTable.stopTrackingPartitions(producingTaskExecutorId);

		for (ResultPartitionID resultPartitionId : resultPartitionIds) {
			internalStopTrackingPartition(resultPartitionId);
		}
	}

	@Override
	public void stopTrackingAndReleasePartitions(Collection resultPartitionIds) {
		Preconditions.checkNotNull(resultPartitionIds);

		// stop tracking partitions to be released and group them by task executor ID
		Map> partitionsToReleaseByResourceId = resultPartitionIds.stream()
			.map(this::internalStopTrackingPartition)
			.filter(Optional::isPresent)
			.map(Optional::get)
			.collect(Collectors.groupingBy(
				partitionMetaData -> partitionMetaData.producingTaskExecutorResourceId,
				Collectors.mapping(
					partitionMetaData -> partitionMetaData.resultPartitionDeploymentDescriptor,
					toList())));

		partitionsToReleaseByResourceId.forEach(this::internalReleasePartitions);
	}

	@Override
	public void stopTrackingPartitions(Collection resultPartitionIds) {
		Preconditions.checkNotNull(resultPartitionIds);

		resultPartitionIds.forEach(this::internalStopTrackingPartition);
	}

	@Override
	public void stopTrackingAndReleasePartitionsFor(ResourceID producingTaskExecutorId) {
		Preconditions.checkNotNull(producingTaskExecutorId);

		// this is a bit icky since we make 2 calls to pT#stopTrackingPartitions
		Collection resultPartitionIds = partitionTable.stopTrackingPartitions(producingTaskExecutorId);

		stopTrackingAndReleasePartitions(resultPartitionIds);
	}

	@Override
	public boolean isTrackingPartitionsFor(ResourceID producingTaskExecutorId) {
		Preconditions.checkNotNull(producingTaskExecutorId);

		return partitionTable.hasTrackedPartitions(producingTaskExecutorId);
	}

	@Override
	public boolean isPartitionTracked(final ResultPartitionID resultPartitionID) {
		Preconditions.checkNotNull(resultPartitionID);

		return partitionInfos.containsKey(resultPartitionID);
	}

	private Optional internalStopTrackingPartition(ResultPartitionID resultPartitionId) {
		final PartitionInfo partitionInfo = partitionInfos.remove(resultPartitionId);
		if (partitionInfo == null) {
			return Optional.empty();
		}
		partitionTable.stopTrackingPartitions(partitionInfo.producingTaskExecutorResourceId, Collections.singletonList(resultPartitionId));

		return Optional.of(partitionInfo);
	}

	private void internalReleasePartitions(
		ResourceID potentialPartitionLocation,
		Collection partitionDeploymentDescriptors) {

		internalReleasePartitionsOnTaskExecutor(potentialPartitionLocation, partitionDeploymentDescriptors);
		internalReleasePartitionsOnShuffleMaster(partitionDeploymentDescriptors);
	}

	private void internalReleasePartitionsOnTaskExecutor(
		ResourceID potentialPartitionLocation,
		Collection partitionDeploymentDescriptors) {

		final List partitionsRequiringRpcReleaseCalls = partitionDeploymentDescriptors.stream()
			.map(ResultPartitionDeploymentDescriptor::getShuffleDescriptor)
			.filter(descriptor -> descriptor.storesLocalResourcesOn().isPresent())
			.map(ShuffleDescriptor::getResultPartitionID)
			.collect(Collectors.toList());

		if (!partitionsRequiringRpcReleaseCalls.isEmpty()) {
			taskExecutorGatewayLookup
				.lookup(potentialPartitionLocation)
				.ifPresent(taskExecutorGateway ->
					taskExecutorGateway.releasePartitions(jobId, partitionsRequiringRpcReleaseCalls));
		}
	}

	private void internalReleasePartitionsOnShuffleMaster(Collection partitionDeploymentDescriptors) {
		partitionDeploymentDescriptors.stream()
			.map(ResultPartitionDeploymentDescriptor::getShuffleDescriptor)
			.forEach(shuffleMaster::releasePartitionExternally);
	}

	private static final class PartitionInfo {
		public final ResourceID producingTaskExecutorResourceId;
		public final ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor;

		private PartitionInfo(ResourceID producingTaskExecutorResourceId, ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor) {
			this.producingTaskExecutorResourceId = producingTaskExecutorResourceId;
			this.resultPartitionDeploymentDescriptor = resultPartitionDeploymentDescriptor;
		}
	}
}