org.apache.flink.runtime.broadcast.BroadcastVariableManager Maven / Gradle / Ivy

Go to download
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.broadcast;

import org.apache.flink.api.common.typeutils.TypeSerializerFactory;
import org.apache.flink.runtime.io.network.api.reader.MutableReader;
import org.apache.flink.runtime.operators.BatchTask;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * The BroadcastVariableManager is used to manage the materialization of broadcast variables. References to materialized
 * broadcast variables are cached and shared between parallel subtasks. A reference count is maintained to track whether
 * the materialization may be cleaned up.
 */
public class BroadcastVariableManager {

	private final ConcurrentHashMap> variables =
							new ConcurrentHashMap>(16);

	// --------------------------------------------------------------------------------------------

	/**
	 * Materializes the broadcast variable for the given name, scoped to the given task and its iteration superstep. An
	 * existing materialization created by another parallel subtask may be returned, if it hasn't expired yet.
	 */
	public  BroadcastVariableMaterialization materializeBroadcastVariable(String name, int superstep, BatchTask holder,
			MutableReader reader, TypeSerializerFactory serializerFactory) throws IOException {
		final BroadcastVariableKey key = new BroadcastVariableKey(holder.getEnvironment().getJobVertexId(), name, superstep);

		while (true) {
			final BroadcastVariableMaterialization newMat = new BroadcastVariableMaterialization(key);

			final BroadcastVariableMaterialization previous = variables.putIfAbsent(key, newMat);

			@SuppressWarnings("unchecked")
			final BroadcastVariableMaterialization materialization = (previous == null) ? newMat : (BroadcastVariableMaterialization) previous;

			try {
				materialization.materializeVariable(reader, serializerFactory, holder);
				return materialization;
			}
			catch (MaterializationExpiredException e) {
				// concurrent release. as an optimization, try to replace the previous one with our version. otherwise we might spin for a while
				// until the releaser removes the variable
				// NOTE: This would also catch a bug prevented an expired materialization from ever being removed, so it acts as a future safeguard

				boolean replaceSuccessful = false;
				try {
					replaceSuccessful = variables.replace(key, materialization, newMat);
				}
				catch (Throwable t) {}

				if (replaceSuccessful) {
					try {
						newMat.materializeVariable(reader, serializerFactory, holder);
						return newMat;
					}
					catch (MaterializationExpiredException ee) {
						// can still happen in cases of extreme races and fast tasks
						// fall through the loop;
					}
				}
				// else fall through the loop
			}
		}
	}

	public void releaseReference(String name, int superstep, BatchTask referenceHolder) {
		BroadcastVariableKey key = new BroadcastVariableKey(referenceHolder.getEnvironment().getJobVertexId(), name, superstep);
		releaseReference(key, referenceHolder);
	}

	public void releaseReference(BroadcastVariableKey key, BatchTask referenceHolder) {
		BroadcastVariableMaterialization mat = variables.get(key);

		// release this reference
		if (mat.decrementReference(referenceHolder)) {
			// remove if no one holds a reference and no one concurrently replaced the entry
			variables.remove(key, mat);
		}
	}

	public void releaseAllReferencesFromTask(BatchTask referenceHolder) {
		// go through all registered variables
		for (Map.Entry> entry : variables.entrySet()) {
			BroadcastVariableMaterialization mat = entry.getValue();

			// release the reference
			if (mat.decrementReferenceIfHeld(referenceHolder)) {
				// remove if no one holds a reference and no one concurrently replaced the entry
				variables.remove(entry.getKey(), mat);
			}
		}
	}

	// --------------------------------------------------------------------------------------------

	public int getNumberOfVariablesWithReferences() {
		return this.variables.size();
	}
}