org.springframework.yarn.am.AppmasterCmTemplate.java Source code

Introduction

Here is the source code for org.springframework.yarn.am.AppmasterCmTemplate.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.am;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.ArrayList;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.api.ContainerManagementProtocol;
import org.apache.hadoop.yarn.api.protocolrecords.GetContainerStatusesRequest;
import org.apache.hadoop.yarn.api.protocolrecords.StartContainersRequest;
import org.apache.hadoop.yarn.api.protocolrecords.StartContainersResponse;
import org.apache.hadoop.yarn.api.protocolrecords.StopContainersRequest;
import org.apache.hadoop.yarn.api.protocolrecords.StopContainersResponse;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.Token;
import org.apache.hadoop.yarn.client.api.NMTokenCache;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.security.NMTokenIdentifier;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.apache.hadoop.yarn.util.Records;
import org.springframework.yarn.rpc.YarnRpcAccessor;
import org.springframework.yarn.rpc.YarnRpcCallback;

/**
 * Template implementation for {@link AppmasterCmOperations} wrapping
 * communication using {@link ContainerManagementProtocol}. Methods for this
 * template wraps possible exceptions into Spring Dao exception hierarchy.
 *
 * @author Janne Valkealahti
 *
 */
public class AppmasterCmTemplate extends YarnRpcAccessor<ContainerManagementProtocol>
        implements AppmasterCmOperations {

    /** Container we're working for */
    private final Container container;

    /**
     * Instantiates a new AppmasterCmTemplate.
     *
     * @param config the hadoop configation
     * @param container the {@link Container}
     */
    public AppmasterCmTemplate(Configuration config, Container container) {
        super(ContainerManagementProtocol.class, config);
        this.container = container;
    }

    @Override
    public StartContainersResponse startContainers(final StartContainersRequest request) {
        return execute(new YarnRpcCallback<StartContainersResponse, ContainerManagementProtocol>() {
            @Override
            public StartContainersResponse doInYarn(ContainerManagementProtocol proxy)
                    throws YarnException, IOException {
                return proxy.startContainers(request);
            }
        });
    }

    @Override
    public StopContainersResponse stopContainers() {
        return execute(new YarnRpcCallback<StopContainersResponse, ContainerManagementProtocol>() {
            @Override
            public StopContainersResponse doInYarn(ContainerManagementProtocol proxy)
                    throws YarnException, IOException {
                StopContainersRequest request = Records.newRecord(StopContainersRequest.class);
                ArrayList<ContainerId> ids = new ArrayList<ContainerId>();
                ids.add(container.getId());
                request.setContainerIds(ids);
                return proxy.stopContainers(request);
            }
        });
    }

    @Override
    public ContainerStatus getContainerStatus() {
        return execute(new YarnRpcCallback<ContainerStatus, ContainerManagementProtocol>() {
            @Override
            public ContainerStatus doInYarn(ContainerManagementProtocol proxy) throws YarnException, IOException {
                GetContainerStatusesRequest request = Records.newRecord(GetContainerStatusesRequest.class);
                ArrayList<ContainerId> ids = new ArrayList<ContainerId>();
                ids.add(container.getId());
                request.setContainerIds(ids);
                return proxy.getContainerStatuses(request).getContainerStatuses().get(0);
            }
        });
    }

    @Override
    protected InetSocketAddress getRpcAddress(Configuration config) {
        String cmIpPortStr = container.getNodeId().getHost() + ":" + container.getNodeId().getPort();
        return NetUtils.createSocketAddr(cmIpPortStr);
    }

    @Override
    protected UserGroupInformation getUser() {
        InetSocketAddress rpcAddress = getRpcAddress(getConfiguration());

        // TODO: at some point remove static cache
        Token token = NMTokenCache.getNMToken(container.getNodeId().toString());

        // this is what node manager requires for auth
        UserGroupInformation user = UserGroupInformation
                .createRemoteUser(container.getId().getApplicationAttemptId().toString());
        org.apache.hadoop.security.token.Token<NMTokenIdentifier> nmToken = ConverterUtils.convertFromYarn(token,
                rpcAddress);
        user.addToken(nmToken);

        return user;
    }

}