co.cask.cdap.internal.app.runtime.batch.MapReduceClassLoader.java Source code

Java tutorial

Introduction

Here is the source code for co.cask.cdap.internal.app.runtime.batch.MapReduceClassLoader.java

Source

/*
 * Copyright  2015-2016 Cask Data, Inc.
 *
 * Licensed 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 co.cask.cdap.internal.app.runtime.batch;

import co.cask.cdap.api.plugin.Plugin;
import co.cask.cdap.common.conf.CConfiguration;
import co.cask.cdap.common.conf.Constants;
import co.cask.cdap.common.io.Locations;
import co.cask.cdap.common.lang.CombineClassLoader;
import co.cask.cdap.common.lang.Delegators;
import co.cask.cdap.common.lang.ProgramClassLoader;
import co.cask.cdap.common.lang.jar.BundleJarUtil;
import co.cask.cdap.common.logging.LoggingContext;
import co.cask.cdap.common.logging.LoggingContextAccessor;
import co.cask.cdap.common.utils.DirUtils;
import co.cask.cdap.internal.app.runtime.batch.distributed.DistributedMapReduceTaskContextProvider;
import co.cask.cdap.internal.app.runtime.batch.distributed.MapReduceContainerLauncher;
import co.cask.cdap.internal.app.runtime.plugin.PluginClassLoaders;
import co.cask.cdap.internal.app.runtime.plugin.PluginInstantiator;
import co.cask.cdap.internal.app.runtime.workflow.WorkflowProgramInfo;
import co.cask.cdap.logging.context.MapReduceLoggingContext;
import co.cask.cdap.logging.context.WorkflowProgramLoggingContext;
import co.cask.cdap.proto.ProgramType;
import co.cask.cdap.proto.id.ProgramId;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.Service;
import com.google.inject.Injector;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.MRJobConfig;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.twill.api.RunId;
import org.apache.twill.filesystem.Location;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;

/**
 * A {@link ClassLoader} for YARN application isolation. Classes from
 * the application JARs are loaded in preference to the parent loader.
 *
 * The delegation order is:
 *
 * ProgramClassLoader -> Plugin Lib ClassLoader -> Plugins Export-Package ClassLoaders -> System ClassLoader
 */
public class MapReduceClassLoader extends CombineClassLoader implements AutoCloseable {

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

    private final Parameters parameters;
    // Supplier for MapReduceTaskContextProvider. Need to wrap it with a supplier to delay calling
    // MapReduceTaskContextProvider.start() since it shouldn't be called in constructor.
    private final Supplier<MapReduceTaskContextProvider> taskContextProviderSupplier;
    private MapReduceTaskContextProvider taskContextProvider;

    /**
     * Finds the {@link MapReduceClassLoader} from the {@link ClassLoader} inside the given {@link Configuration}.
     *
     * @throws IllegalArgumentException if no {@link MapReduceClassLoader} can be found from the {@link Configuration}.
     */
    public static MapReduceClassLoader getFromConfiguration(Configuration configuration) {
        return Delegators.getDelegate(configuration.getClassLoader(), MapReduceClassLoader.class);
    }

    /**
     * Constructor. It creates classloader for MapReduce from information
     * gathered through {@link MapReduceContextConfig}. This method is called by {@link MapReduceContainerLauncher}.
     */
    @SuppressWarnings("unused")
    public MapReduceClassLoader() {
        this(new Parameters(), new TaskContextProviderFactory() {
            @Override
            public MapReduceTaskContextProvider create(CConfiguration cConf, Configuration hConf) {
                Preconditions.checkState(!MapReduceTaskContextProvider.isLocal(hConf),
                        "Expected to be in distributed mode.");
                return new DistributedMapReduceTaskContextProvider(cConf, hConf);
            }
        });
    }

    /**
     * Constructs a ClassLoader that load classes from the programClassLoader, then from the plugin lib ClassLoader,
     * followed by plugin Export-Package ClassLoader and with the system ClassLoader last.
     * This constructor should only be called from {@link MapReduceRuntimeService} only.
     */
    MapReduceClassLoader(final Injector injector, CConfiguration cConf, Configuration hConf,
            ClassLoader programClassLoader, Map<String, Plugin> plugins,
            @Nullable PluginInstantiator pluginInstantiator) {
        this(new Parameters(cConf, hConf, programClassLoader, plugins, pluginInstantiator),
                new TaskContextProviderFactory() {
                    @Override
                    public MapReduceTaskContextProvider create(CConfiguration cConf, Configuration hConf) {
                        return new MapReduceTaskContextProvider(injector);
                    }
                });
    }

    /**
     * Constructs a ClassLoader based on the given {@link Parameters} and also uses the given
     * {@link TaskContextProviderFactory} to create {@link MapReduceTaskContextProvider} on demand.
     */
    private MapReduceClassLoader(final Parameters parameters,
            final TaskContextProviderFactory contextProviderFactory) {
        super(null, createDelegates(parameters));
        this.parameters = parameters;
        this.taskContextProviderSupplier = new Supplier<MapReduceTaskContextProvider>() {
            @Override
            public MapReduceTaskContextProvider get() {
                return contextProviderFactory.create(parameters.getCConf(), parameters.getHConf());
            }
        };
    }

    /**
     * Returns the {@link MapReduceTaskContextProvider} associated with this ClassLoader.
     */
    public MapReduceTaskContextProvider getTaskContextProvider() {
        // Logging context needs to be set in main thread.
        LoggingContext loggingContext = createMapReduceLoggingContext();
        LoggingContextAccessor.setLoggingContext(loggingContext);

        synchronized (this) {
            taskContextProvider = Optional.fromNullable(taskContextProvider).or(taskContextProviderSupplier);
        }
        taskContextProvider.startAndWait();
        return taskContextProvider;
    }

    /**
     * Creates logging context for MapReduce program. If the program is started
     * by Workflow an instance of {@link WorkflowProgramLoggingContext} is returned,
     * otherwise an instance of {@link MapReduceLoggingContext} is returned.
     */
    private LoggingContext createMapReduceLoggingContext() {
        MapReduceContextConfig contextConfig = new MapReduceContextConfig(parameters.getHConf());
        ProgramId programId = contextConfig.getProgramId();
        RunId runId = contextConfig.getRunId();
        WorkflowProgramInfo workflowProgramInfo = contextConfig.getWorkflowProgramInfo();
        if (workflowProgramInfo == null) {
            return new MapReduceLoggingContext(programId.getNamespace(), programId.getApplication(),
                    programId.getProgram(), runId.getId());
        }
        String workflowId = workflowProgramInfo.getName();
        String workflowRunId = workflowProgramInfo.getRunId().getId();
        return new WorkflowProgramLoggingContext(programId.getNamespace(), programId.getApplication(), workflowId,
                workflowRunId, ProgramType.MAPREDUCE, programId.getProgram());
    }

    /**
     * Returns the program {@link ProgramClassLoader} used to construct this ClassLoader.
     */
    public ClassLoader getProgramClassLoader() {
        return parameters.getProgramClassLoader();
    }

    /**
     * Returns the {@link PluginInstantiator} associated with this ClassLoader.
     */
    @Nullable
    public PluginInstantiator getPluginInstantiator() {
        return parameters.getPluginInstantiator();
    }

    @Override
    public void close() throws Exception {
        try {
            MapReduceTaskContextProvider provider;
            synchronized (this) {
                provider = taskContextProvider;
            }
            if (provider != null) {
                Service.State state = provider.state();
                if (state == Service.State.STARTING || state == Service.State.RUNNING) {
                    provider.stopAndWait();
                }
            }
        } catch (Exception e) {
            // This is non-fatal, since the container is already done.
            LOG.warn("Exception while stopping MapReduceTaskContextProvider", e);
        }
    }

    /**
     * Creates the delegating list of ClassLoader.
     */
    private static List<ClassLoader> createDelegates(Parameters parameters) {
        return ImmutableList.of(parameters.getProgramClassLoader(), parameters.getFilteredPluginsClassLoader(),
                MapReduceClassLoader.class.getClassLoader());
    }

    /**
     * A container class for holding parameters for the construction of the MapReduceClassLoader.
     * It is needed because we need all parameters available when calling super constructor.
     */
    private static final class Parameters {

        private final CConfiguration cConf;
        private final Configuration hConf;
        private final ClassLoader programClassLoader;
        private final PluginInstantiator pluginInstantiator;
        private final ClassLoader filteredPluginsClassLoader;

        /**
         * Creates from the Job Configuration
         */
        Parameters() {
            this(createContextConfig());
        }

        Parameters(MapReduceContextConfig contextConfig) {
            this(contextConfig, createProgramClassLoader(contextConfig));
        }

        Parameters(MapReduceContextConfig contextConfig, ClassLoader programClassLoader) {
            this(contextConfig.getCConf(), contextConfig.getHConf(), programClassLoader, contextConfig.getPlugins(),
                    createPluginInstantiator(contextConfig, programClassLoader));
        }

        /**
         * Creates from the given ProgramClassLoader with plugin classloading support.
         */
        Parameters(CConfiguration cConf, Configuration hConf, ClassLoader programClassLoader,
                Map<String, Plugin> plugins, @Nullable PluginInstantiator pluginInstantiator) {
            this.cConf = cConf;
            this.hConf = hConf;
            this.programClassLoader = programClassLoader;
            this.pluginInstantiator = pluginInstantiator;
            this.filteredPluginsClassLoader = PluginClassLoaders.createFilteredPluginsClassLoader(plugins,
                    pluginInstantiator);
        }

        ClassLoader getProgramClassLoader() {
            return programClassLoader;
        }

        PluginInstantiator getPluginInstantiator() {
            return pluginInstantiator;
        }

        ClassLoader getFilteredPluginsClassLoader() {
            return filteredPluginsClassLoader;
        }

        CConfiguration getCConf() {
            return cConf;
        }

        Configuration getHConf() {
            return hConf;
        }

        private static MapReduceContextConfig createContextConfig() {
            Configuration conf = new Configuration(new YarnConfiguration());
            conf.addResource(new Path(MRJobConfig.JOB_CONF_FILE));
            return new MapReduceContextConfig(conf);
        }

        /**
         * Creates a program {@link ClassLoader} based on the MR job config.
         */
        private static ClassLoader createProgramClassLoader(MapReduceContextConfig contextConfig) {
            // In distributed mode, the program is created by expanding the program jar.
            // The program jar is localized to container with the program jar name.
            // It's ok to expand to a temp dir in local directory, as the YARN container will be gone.
            Location programLocation = Locations.toLocation(new File(contextConfig.getProgramJarName()));
            try {
                File unpackDir = DirUtils.createTempDir(new File(System.getProperty("user.dir")));
                LOG.info("Create ProgramClassLoader from {}, expand to {}", programLocation, unpackDir);

                BundleJarUtil.unJar(programLocation, unpackDir);
                return ProgramClassLoader.create(contextConfig.getCConf(), unpackDir,
                        contextConfig.getHConf().getClassLoader());
            } catch (IOException e) {
                LOG.error("Failed to create ProgramClassLoader", e);
                throw Throwables.propagate(e);
            }
        }

        /**
         * Returns a new {@link PluginInstantiator} or {@code null} if no plugin is supported.
         */
        @Nullable
        private static PluginInstantiator createPluginInstantiator(MapReduceContextConfig contextConfig,
                ClassLoader programClassLoader) {
            String pluginArchive = contextConfig.getHConf().get(Constants.Plugin.ARCHIVE);
            if (pluginArchive == null) {
                return null;
            }
            return new PluginInstantiator(contextConfig.getCConf(), programClassLoader, new File(pluginArchive));
        }
    }

    /**
     * A private interface to help abstract out which type of {@link MapReduceTaskContextProvider} is created,
     * depending on the runtime environment.
     */
    private interface TaskContextProviderFactory {

        /**
         * Returns a new instance of {@link MapReduceTaskContextProvider}.
         */
        MapReduceTaskContextProvider create(CConfiguration cConf, Configuration hConf);
    }
}