org.apache.storm.scheduler.utils.ArtifactoryConfigLoader.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.storm.scheduler.utils.ArtifactoryConfigLoader.java

Source

/**
 * 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.storm.scheduler.utils;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Collections;
import java.util.Comparator;
import java.util.Map;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.client.HttpClient;
import org.apache.http.client.ResponseHandler;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.utils.URIBuilder;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.util.EntityUtils;
import org.apache.storm.Config;
import org.apache.storm.DaemonConfig;
import org.apache.storm.utils.ConfigUtils;
import org.apache.storm.utils.ServerConfigUtils;
import org.apache.storm.utils.Time;
import org.apache.storm.utils.Utils;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;
import org.json.simple.parser.ParseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.constructor.SafeConstructor;

/**
 * A dynamic loader that can load scheduler configurations for user resource guarantees from Artifactory (an artifact repository manager).
 */
public class ArtifactoryConfigLoader implements IConfigLoader {
    protected static final String LOCAL_ARTIFACT_DIR = "scheduler_artifacts";
    static final String cacheFilename = "latest.yaml";
    private static final String DEFAULT_ARTIFACTORY_BASE_DIRECTORY = "/artifactory";
    private static final int DEFAULT_POLLTIME_SECS = 600;
    private static final int DEFAULT_TIMEOUT_SECS = 10;
    private static final String ARTIFACTORY_SCHEME_PREFIX = "artifactory+";

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

    private Map<String, Object> conf;
    private int artifactoryPollTimeSecs = DEFAULT_POLLTIME_SECS;
    private boolean cacheInitialized = false;
    // Location of the file in the artifactory archive.  Also used to name file in cache.
    private String localCacheDir;
    private String baseDirectory = DEFAULT_ARTIFACTORY_BASE_DIRECTORY;
    private int lastReturnedTime = 0;
    private int timeoutSeconds = DEFAULT_TIMEOUT_SECS;
    private Map lastReturnedValue;
    private URI targetURI = null;
    private JSONParser jsonParser;
    private String scheme;

    public ArtifactoryConfigLoader(Map<String, Object> conf) {
        this.conf = conf;
        Integer thisTimeout = (Integer) conf.get(DaemonConfig.SCHEDULER_CONFIG_LOADER_TIMEOUT_SECS);
        if (thisTimeout != null) {
            timeoutSeconds = thisTimeout;
        }
        Integer thisPollTime = (Integer) conf.get(DaemonConfig.SCHEDULER_CONFIG_LOADER_POLLTIME_SECS);
        if (thisPollTime != null) {
            artifactoryPollTimeSecs = thisPollTime;
        }
        String thisBase = (String) conf.get(DaemonConfig.SCHEDULER_CONFIG_LOADER_ARTIFACTORY_BASE_DIRECTORY);
        if (thisBase != null) {
            baseDirectory = thisBase;
        }
        String uriString = (String) conf.get(DaemonConfig.SCHEDULER_CONFIG_LOADER_URI);
        if (uriString == null) {
            LOG.error("No URI defined in {} configuration.", DaemonConfig.SCHEDULER_CONFIG_LOADER_URI);
        } else {
            try {
                targetURI = new URI(uriString);
                scheme = targetURI.getScheme().substring(ARTIFACTORY_SCHEME_PREFIX.length());
            } catch (URISyntaxException e) {
                LOG.error("Failed to parse uri={}", uriString);
            }
        }
        jsonParser = new JSONParser();
    }

    /**
     * Load the configs associated with the configKey from the targetURI.
     * @param configKey The key from which we want to get the scheduler config.
     * @return The scheduler configuration if exists; null otherwise.
     */
    @Override
    public Map load(String configKey) {
        if (targetURI == null) {
            return null;
        }

        // Check for new file every so often
        int currentTimeSecs = Time.currentTimeSecs();
        if (lastReturnedValue != null && ((currentTimeSecs - lastReturnedTime) < artifactoryPollTimeSecs)) {
            LOG.debug(
                    "currentTimeSecs: {}; lastReturnedTime {}; artifactoryPollTimeSecs: {}. Returning our last map.",
                    currentTimeSecs, lastReturnedTime, artifactoryPollTimeSecs);
            return (Map) lastReturnedValue.get(configKey);
        }

        try {
            Map raw = loadFromURI(targetURI);
            if (raw != null) {
                return (Map) raw.get(configKey);
            }
        } catch (Exception e) {
            LOG.error("Failed to load from uri {}", targetURI);
        }
        return null;
    }

    /**
     * A private class used to check the response coming back from httpclient.
     */
    private static class GETStringResponseHandler implements ResponseHandler<String> {
        private static GETStringResponseHandler singleton = null;

        /**
         * @return a singleton httpclient GET response handler
         */
        public static GETStringResponseHandler getInstance() {
            if (singleton == null) {
                singleton = new GETStringResponseHandler();
            }
            return singleton;
        }

        /**
         * @param response The http response to verify.
         * @return null on failure or the response string if return code is in 200 range
         */
        @Override
        public String handleResponse(final HttpResponse response) throws IOException {
            int status = response.getStatusLine().getStatusCode();
            HttpEntity entity = response.getEntity();
            String entityString = (entity != null ? EntityUtils.toString(entity) : null);
            if (status >= 200 && status < 300) {
                return entityString;
            } else {
                LOG.error("Got unexpected response code {}; entity: {}", status, entityString);
                return null;
            }
        }
    }

    /**
     * @param api null if we are trying to download artifact, otherwise a string to call REST api,
     *        e.g. "/api/storage"
     * @param artifact location of artifact
     * @param host Artifactory hostname
     * @param port Artifactory port
     * @return null on failure or the response string if return code is in 200 range
     *
     * <p>Protected so we can override this in unit tests
     */
    protected String doGet(String api, String artifact, String host, Integer port) {
        URIBuilder builder = new URIBuilder().setScheme(scheme).setHost(host).setPort(port);

        String path = null;
        if (api != null) {
            path = baseDirectory + "/" + api + "/" + artifact;
        } else {
            path = baseDirectory + "/" + artifact;
        }

        // Get rid of multiple '/' in url
        path = path.replaceAll("/[/]+", "/");
        builder.setPath(path);

        RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(timeoutSeconds * 1000).build();
        HttpClient httpclient = HttpClientBuilder.create().setDefaultRequestConfig(requestConfig).build();

        String returnValue;
        try {
            LOG.debug("About to issue a GET to {}", builder);
            HttpGet httpget = new HttpGet(builder.build());
            String responseBody;
            responseBody = httpclient.execute(httpget, GETStringResponseHandler.getInstance());
            returnValue = responseBody;
        } catch (Exception e) {
            LOG.error("Received exception while connecting to Artifactory", e);
            returnValue = null;
        }

        LOG.debug("Returning {}", returnValue);
        return returnValue;
    }

    private JSONObject getArtifactMetadata(String location, String host, Integer port) {
        String metadataStr = null;

        metadataStr = doGet("/api/storage", location, host, port);

        if (metadataStr == null) {
            return null;
        }

        JSONObject returnValue;
        try {
            returnValue = (JSONObject) jsonParser.parse(metadataStr);
        } catch (ParseException e) {
            LOG.error("Could not parse JSON string {}", metadataStr, e);
            return null;
        }

        return returnValue;
    }

    private class DirEntryCompare implements Comparator<JSONObject> {
        @Override
        public int compare(JSONObject o1, JSONObject o2) {
            return ((String) o1.get("uri")).compareTo((String) o2.get("uri"));
        }
    }

    private String loadMostRecentArtifact(String location, String host, Integer port) {
        // Is this a directory or is it a file?
        JSONObject json = getArtifactMetadata(location, host, port);
        if (json == null) {
            LOG.error("got null metadata");
            return null;
        }
        String downloadURI = (String) json.get("downloadUri");

        // This means we are pointing at a file.
        if (downloadURI != null) {
            // Then get it and return the file as string.
            String returnValue = doGet(null, location, host, port);
            saveInArtifactoryCache(returnValue);
            return returnValue;
        }

        // This should mean that we were pointed at a directory.  
        // Find the most recent child and load that.
        JSONArray msg = (JSONArray) json.get("children");
        if (msg == null || msg.size() == 0) {
            LOG.error("Expected directory children not present");
            return null;
        }
        JSONObject newest = (JSONObject) Collections.max(msg, new DirEntryCompare());
        if (newest == null) {
            LOG.error("Failed to find most recent artifact uri in {}", location);
            return null;
        }

        String uri = (String) newest.get("uri");
        if (uri == null) {
            LOG.error("Expected directory uri not present");
            return null;
        }
        String returnValue = doGet(null, location + uri, host, port);
        saveInArtifactoryCache(returnValue);
        return returnValue;
    }

    private void updateLastReturned(Map ret) {
        lastReturnedTime = Time.currentTimeSecs();
        lastReturnedValue = ret;
    }

    private Map loadFromFile(File file) {
        Map ret = null;

        try {
            ret = (Map) Utils.readYamlFile(file.getCanonicalPath());
        } catch (IOException e) {
            LOG.error("Filed to load from file. Exception: {}", e.getMessage());
        }

        if (ret != null) {
            try {
                LOG.debug("returning a new map from file {}", file.getCanonicalPath());
            } catch (java.io.IOException e) {
                LOG.debug("Could not get PATH from file object in debug print. Ignoring");
            }
            return ret;
        }

        return null;
    }

    private Map getLatestFromCache() {
        String localFileName = localCacheDir + File.separator + cacheFilename;
        return loadFromFile(new File(localFileName));
    }

    private void saveInArtifactoryCache(String yamlData) {
        if (yamlData == null) {
            LOG.warn("Will not save null data into the artifactory cache");
            return;
        }

        String localFileName = localCacheDir + File.separator + cacheFilename;

        File cacheFile = new File(localFileName);
        try (FileOutputStream fos = new FileOutputStream(cacheFile)) {
            fos.write(yamlData.getBytes());
            fos.flush();
        } catch (IOException e) {
            LOG.error("Received exception when writing file {}.  Attempting delete", localFileName, e);
            try {
                cacheFile.delete();
            } catch (Exception deleteException) {
                LOG.error("Received exception when deleting file {}.", localFileName, deleteException);
            }
        }
    }

    private void makeArtifactoryCache(String location) throws IOException {
        // First make the cache dir
        String localDirName = ServerConfigUtils.masterLocalDir(conf) + File.separator + LOCAL_ARTIFACT_DIR;
        File dir = new File(localDirName);
        if (!dir.exists()) {
            dir.mkdirs();
        }

        localCacheDir = localDirName + File.separator + location.replaceAll(File.separator, "_");
        dir = new File(localCacheDir);
        if (!dir.exists()) {
            dir.mkdir();
        }
        cacheInitialized = true;
    }

    private Map loadFromURI(URI uri) throws IOException {
        String host = uri.getHost();
        Integer port = uri.getPort();
        String location = uri.getPath();
        if (location.toLowerCase().startsWith(baseDirectory.toLowerCase())) {
            location = location.substring(baseDirectory.length());
        }

        if (!cacheInitialized) {
            makeArtifactoryCache(location);
        }

        // Get the most recent artifact as a String, and then parse the yaml
        String yamlConfig = loadMostRecentArtifact(location, host, port);

        // If we failed to get anything from Artifactory try to get it from our local cache
        if (yamlConfig == null) {
            Map ret = getLatestFromCache();
            updateLastReturned(ret);
            return ret;
        }

        // Now parse it and return the map.
        Yaml yaml = new Yaml(new SafeConstructor());
        Map ret = null;
        try {
            ret = (Map) yaml.load(yamlConfig);
        } catch (Exception e) {
            LOG.error("Could not parse yaml.");
            return null;
        }

        if (ret != null) {
            LOG.debug("returning a new map from Artifactory");
            updateLastReturned(ret);
            return ret;
        }

        return null;
    }
}