org.springframework.xd.yarn.XdYarnAdminInfoCommand.java Source code

Introduction

Here is the source code for org.springframework.xd.yarn.XdYarnAdminInfoCommand.java

Source

/*
 * Copyright 2015 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.xd.yarn;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Properties;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

import org.springframework.beans.factory.config.YamlPropertiesFactoryBean;
import org.springframework.boot.cli.command.AbstractCommand;
import org.springframework.boot.cli.command.status.ExitStatus;
import org.springframework.boot.cli.util.Log;
import org.springframework.core.io.FileSystemResource;

/**
 * A custom command for listing the http ports used by admin servers.
 *
 * @author Thomas Risberg
 */
@SuppressWarnings("unchecked")
public class XdYarnAdminInfoCommand extends AbstractCommand {

    private static final String XD_YARN_ADMIN_INFO_COMMAND_NAME = "admininfo";
    private static final String XD_YARN_ADMIN_INFO_COMMAND_DESC = "List http host and port info for the admin servers";

    public static final String ZK_NAMESPACE = "zk.namespace";
    public static final String ZK_CLIENT_CONNECT = "zk.client.connect";

    public static final String XD_ADMINS = "admins";

    private String zkNamespace = "xd";
    private String zkClientConnect = "localhost:2181";

    private ZooKeeper zk;

    /**
     * Call super constructor for new {@link org.springframework.boot.cli.command.AbstractCommand} instance.
     */
    public XdYarnAdminInfoCommand() {
        super(XD_YARN_ADMIN_INFO_COMMAND_NAME, XD_YARN_ADMIN_INFO_COMMAND_DESC);
    }

    @Override
    public ExitStatus run(String... args) throws InterruptedException {
        try {
            init();
            Log.info("Admins: " + getAdmins());
        } catch (IOException e) {
            Log.error("IOException accessing admin info: " + e.getMessage());
        } catch (KeeperException e) {
            Log.error("KeeperException accessing admin info: " + e.getMessage());
        } finally {
            close();
        }
        return ExitStatus.OK;
    }

    private void init() throws IOException {
        YamlPropertiesFactoryBean ypfb = new YamlPropertiesFactoryBean();
        ypfb.setResources(new FileSystemResource(System.getProperty("spring.config.location")));
        Properties conf = ypfb.getObject();
        if (conf.containsKey(ZK_NAMESPACE)) {
            zkNamespace = (String) conf.get(ZK_NAMESPACE);
        }
        if (conf.containsKey(ZK_CLIENT_CONNECT)) {
            zkClientConnect = (String) conf.get(ZK_CLIENT_CONNECT);
        }
        zk = new ZooKeeper(zkClientConnect, 5000, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
            }
        });
    }

    private void close() throws InterruptedException {
        if (zk != null) {
            zk.close();
        }
    }

    private List<String> getAdmins() throws KeeperException {
        String adminsPath = "/" + zkNamespace + "/" + XD_ADMINS;
        List<String> results = new ArrayList<String>();
        try {
            Stat stat = zk.exists(adminsPath, false);
            if (stat != null) {
                List<String> admins = zk.getChildren(adminsPath, false);
                for (String admin : admins) {
                    byte[] data = zk.getData(adminsPath + "/" + admin, false, null);
                    try {
                        HashMap<String, Object> dataValues = new ObjectMapper().readValue(new String(data),
                                HashMap.class);
                        results.add("http://" + dataValues.get("host") + ":" + dataValues.get("port"));
                    } catch (IOException e) {
                        Log.error("IOException while retrieving data for " + admin + ": " + e.getMessage());
                    }
                }
            }
        } catch (InterruptedException ignore) {
        }
        return results;
    }
}