org.springframework.yarn.container.CommandLineContainerRunner.java Source code

Introduction

Here is the source code for org.springframework.yarn.container.CommandLineContainerRunner.java

Source

/*
 * Copyright 2013 the original author or authors.
 *
 * 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 org.springframework.yarn.container;

import java.util.List;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.CountDownLatch;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.util.StringUtils;
import org.springframework.yarn.YarnSystemConstants;
import org.springframework.yarn.launch.AbstractCommandLineRunner;
import org.springframework.yarn.launch.ExitStatus;
import org.springframework.yarn.listener.ContainerStateListener;
import org.springframework.yarn.listener.ContainerStateListener.ContainerState;

/**
 * A simple container runner executing a bean
 * named "yarnContainer".
 *
 * @author Janne Valkealahti
 *
 */
public class CommandLineContainerRunner extends AbstractCommandLineRunner<YarnContainer> {

    private static final Log log = LogFactory.getLog(CommandLineContainerRunner.class);

    /** Latch to wait container complete state */
    private CountDownLatch latch;

    private final StateWrapper stateWrapper = new StateWrapper();

    @Override
    protected ExitStatus handleBeanRun(YarnContainer bean, String[] parameters, Set<String> opts) {
        Properties properties = StringUtils.splitArrayElementsIntoProperties(parameters, "=");
        bean.setParameters(properties != null ? properties : new Properties());
        bean.setEnvironment(System.getenv());

        if (log.isDebugEnabled()) {
            log.debug("Starting YarnClient bean: " + StringUtils.arrayToCommaDelimitedString(parameters));
        }

        // use latch if container wants to be long running
        if (bean instanceof LongRunningYarnContainer && ((LongRunningYarnContainer) bean).isWaitCompleteState()) {
            latch = new CountDownLatch(1);
            ((LongRunningYarnContainer) bean).addContainerStateListener(new ContainerStateListener() {
                @Override
                public void state(ContainerState state, Object exit) {
                    stateWrapper.state = state;
                    // TODO: should handle exit value
                    latch.countDown();
                }
            });
        }

        bean.run();

        if (latch != null) {
            try {
                // TODO: should we use timeout?
                latch.await();
            } catch (InterruptedException e) {
                log.debug("Latch interrupted");
            }
        }

        if (log.isDebugEnabled()) {
            log.debug("YarnClient bean complete");
        }

        if (stateWrapper.state != null && stateWrapper.state == ContainerState.FAILED) {
            return ExitStatus.FAILED;
        } else {
            return ExitStatus.COMPLETED;
        }
    }

    @Override
    protected String getDefaultBeanIdentifier() {
        return YarnSystemConstants.DEFAULT_ID_CONTAINER;
    }

    @Override
    protected List<String> getValidOpts() {
        return null;
    }

    public static void main(String[] args) {
        new CommandLineContainerRunner().doMain(args);
    }

    private static class StateWrapper {
        ContainerState state;
    }

}