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

/*
 * 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.annotation.VisibleForTesting;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;

import org.apache.flink.runtime.taskexecutor.ResultPartitionReport;
import org.apache.flink.runtime.taskexecutor.ResultPartitionStatus;
import org.apache.flink.shaded.guava18.com.google.common.collect.HashBasedTable;
import org.apache.flink.shaded.guava18.com.google.common.collect.ImmutableList;
import org.apache.flink.shaded.guava18.com.google.common.collect.Table;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static org.apache.flink.util.Preconditions.checkState;

/**
 * The result partition manager keeps track of all currently produced/consumed partitions of a
 * task manager.
 */
public class ResultPartitionManager implements ResultPartitionProvider {

	private static final Logger LOG = LoggerFactory.getLogger(ResultPartitionManager.class);

	public final Table
			registeredPartitions = HashBasedTable.create();

	private boolean isShutdown;

	public void registerResultPartition(InternalResultPartition partition) throws IOException {
		synchronized (registeredPartitions) {
			checkState(!isShutdown, "Result partition manager already shut down.");

			ResultPartitionID partitionId = partition.getPartitionId();

			InternalResultPartition previous = registeredPartitions.put(
					partitionId.getProducerId(), partitionId.getPartitionId(), partition);

			if (previous != null) {
				throw new IllegalStateException("Result partition already registered.");
			}

			LOG.debug("Registered {}.", partition);
		}
	}

	@Override
	public ResultSubpartitionView createSubpartitionView(
			ResultPartitionID partitionId,
			int subpartitionIndex,
			BufferAvailabilityListener availabilityListener) throws IOException {

		synchronized (registeredPartitions) {
			final InternalResultPartition partition = registeredPartitions.get(partitionId.getProducerId(),
					partitionId.getPartitionId());

			if (partition == null) {
				throw new PartitionNotFoundException(partitionId);
			}

			LOG.debug("Requesting subpartition {} of {}.", subpartitionIndex, partition);

			return partition.createSubpartitionView(subpartitionIndex, availabilityListener);
		}
	}

	public void releasePartitionsProducedBy(ExecutionAttemptID executionId) {
		releasePartitionsProducedBy(executionId, null);
	}

	public void releasePartitionsProducedBy(ExecutionAttemptID executionId, Throwable cause) {
		synchronized (registeredPartitions) {
			final Map partitions =
					registeredPartitions.row(executionId);

			for (InternalResultPartition partition : partitions.values()) {
				partition.release(cause);
			}

			for (IntermediateResultPartitionID partitionId : ImmutableList
					.copyOf(partitions.keySet())) {

				registeredPartitions.remove(executionId, partitionId);
			}

			LOG.debug("Released all partitions produced by {}.", executionId);
		}
	}

	public void releasePartitionsProducedBy(JobID jobId) {
		List matchedResultPartitions = new ArrayList<>();

		synchronized (registeredPartitions) {
			for (InternalResultPartition partition : registeredPartitions.values()) {
				if (partition.getJobId().equals(jobId)) {
					matchedResultPartitions.add(partition);
				}
			}

			for (InternalResultPartition partition : matchedResultPartitions) {
				partition.release();

				ResultPartitionID partitionId = partition.getPartitionId();
				registeredPartitions.remove(partitionId.getProducerId(), partitionId.getPartitionId());
			}

			LOG.info("Released all partitions produced by job {}.", jobId);
		}
	}

	public void shutdown() {
		synchronized (registeredPartitions) {

			LOG.debug("Releasing {} partitions because of shutdown.",
					registeredPartitions.values().size());

			for (InternalResultPartition partition : registeredPartitions.values()) {
				partition.release();
			}

			registeredPartitions.clear();

			isShutdown = true;

			LOG.debug("Successful shutdown.");
		}
	}

	public ResultPartitionReport createResultPartitionReport() {
		List resultPartitionStatuses = new ArrayList<>();

		synchronized (registeredPartitions) {
			for (InternalResultPartition partition : registeredPartitions.values()) {

				if (partition.isFinished) {
					resultPartitionStatuses.add(new ResultPartitionStatus(partition.getJobId(), partition.getPartitionId()));
				}
			}
		}

		return new ResultPartitionReport(resultPartitionStatuses);
	}

	// ------------------------------------------------------------------------
	// Notifications
	// ------------------------------------------------------------------------

	void onConsumedPartition(InternalResultPartition partition) {
		LOG.debug("Received consume notification from {}.", partition);

		if (partition.getPartitionType() == ResultPartitionType.PIPELINED) {
			final InternalResultPartition previous;

			synchronized (registeredPartitions) {
				ResultPartitionID partitionId = partition.getPartitionId();

				previous = registeredPartitions.remove(partitionId.getProducerId(),
					partitionId.getPartitionId());
			}

			// Release the partition if it was successfully removed
			if (partition == previous) {
				partition.release();

				LOG.debug("Released {}.", partition);
			}
		}
	}

	/**
	 * Check if there is no pending partitions. It will check for both pipeline and blocking result partitions,
	 * since the task manager can not exit with both pipeline and blocking result partitions.
	 */
	public boolean areAllPartitionsReleased() {
		synchronized (registeredPartitions) {
			for (ResultPartition partition : registeredPartitions.values()) {
				if (!partition.isReleased()) {
					return false;
				}
			}
			return true;
		}
	}

	@VisibleForTesting
	public boolean containsResultPartitionOfJob(JobID jobId) {
		synchronized (registeredPartitions) {
			for (InternalResultPartition partition : registeredPartitions.values()) {
				if (partition.getJobId().equals(jobId)) {
					return true;
				}
			}
		}

		return false;
	}
}