com.spotify.helios.client.HeliosClient.java Source code

Java tutorial

Introduction

Here is the source code for com.spotify.helios.client.HeliosClient.java

Source

/*
 * Copyright (c) 2014 Spotify AB.
 *
 * 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 com.spotify.helios.client;

import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.escape.Escaper;
import com.google.common.net.UrlEscapers;
import com.google.common.util.concurrent.AsyncFunction;
import com.google.common.util.concurrent.FutureFallback;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningScheduledExecutorService;
import com.google.common.util.concurrent.MoreExecutors;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JavaType;
import com.fasterxml.jackson.databind.type.TypeFactory;
import com.spotify.helios.common.HeliosException;
import com.spotify.helios.common.Json;
import com.spotify.helios.common.Resolver;
import com.spotify.helios.common.SystemClock;
import com.spotify.helios.common.Version;
import com.spotify.helios.common.VersionCompatibility;
import com.spotify.helios.common.descriptors.Deployment;
import com.spotify.helios.common.descriptors.DeploymentGroup;
import com.spotify.helios.common.descriptors.HostStatus;
import com.spotify.helios.common.descriptors.Job;
import com.spotify.helios.common.descriptors.JobId;
import com.spotify.helios.common.descriptors.JobStatus;
import com.spotify.helios.common.descriptors.RolloutOptions;
import com.spotify.helios.common.protocol.CreateDeploymentGroupResponse;
import com.spotify.helios.common.protocol.CreateJobResponse;
import com.spotify.helios.common.protocol.DeploymentGroupStatusResponse;
import com.spotify.helios.common.protocol.HostDeregisterResponse;
import com.spotify.helios.common.protocol.JobDeleteResponse;
import com.spotify.helios.common.protocol.JobDeployResponse;
import com.spotify.helios.common.protocol.JobUndeployResponse;
import com.spotify.helios.common.protocol.RemoveDeploymentGroupResponse;
import com.spotify.helios.common.protocol.RollingUpdateRequest;
import com.spotify.helios.common.protocol.RollingUpdateResponse;
import com.spotify.helios.common.protocol.SetGoalResponse;
import com.spotify.helios.common.protocol.TaskStatusEvents;
import com.spotify.helios.common.protocol.VersionResponse;
import com.spotify.sshagentproxy.AgentProxies;
import com.spotify.sshagentproxy.AgentProxy;

import org.apache.http.client.utils.URIBuilder;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.util.concurrent.Futures.immediateFuture;
import static com.google.common.util.concurrent.Futures.transform;
import static com.google.common.util.concurrent.Futures.withFallback;
import static com.spotify.helios.common.VersionCompatibility.HELIOS_SERVER_VERSION_HEADER;
import static com.spotify.helios.common.VersionCompatibility.HELIOS_VERSION_STATUS_HEADER;
import static java.lang.String.format;
import static java.net.HttpURLConnection.HTTP_BAD_METHOD;
import static java.net.HttpURLConnection.HTTP_BAD_REQUEST;
import static java.net.HttpURLConnection.HTTP_FORBIDDEN;
import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
import static java.net.HttpURLConnection.HTTP_OK;
import static java.util.Arrays.asList;
import static java.util.Collections.singletonList;
import static java.util.concurrent.Executors.newScheduledThreadPool;
import static java.util.concurrent.TimeUnit.SECONDS;

public class HeliosClient implements AutoCloseable {

    private static final Logger log = LoggerFactory.getLogger(HeliosClient.class);

    private final String user;
    private final RequestDispatcher dispatcher;
    private final AtomicBoolean versionWarningLogged = new AtomicBoolean();

    public HeliosClient(final String user, final RequestDispatcher dispatcher) {
        this.user = checkNotNull(user);
        this.dispatcher = checkNotNull(dispatcher);
    }

    @Override
    public void close() throws Exception {
        dispatcher.close();
    }

    private URI uri(final String path) {
        return uri(path, Collections.<String, String>emptyMap());
    }

    private URI uri(final String path, final Map<String, String> query) {
        checkArgument(path.startsWith("/"));

        final URIBuilder builder = new URIBuilder().setScheme("http").setHost("helios").setPath(path);

        for (final Map.Entry<String, String> q : query.entrySet()) {
            builder.addParameter(q.getKey(), q.getValue());
        }
        builder.addParameter("user", user);

        try {
            return builder.build();
        } catch (URISyntaxException e) {
            throw Throwables.propagate(e);
        }
    }

    private String path(final String resource, final Object... params) {
        final String path;
        final Escaper escaper = UrlEscapers.urlPathSegmentEscaper();
        if (params.length == 0) {
            path = resource;
        } else {
            final List<String> encodedParams = Lists.newArrayList();
            for (final Object param : params) {
                encodedParams.add(escaper.escape(param.toString()));
            }
            path = format(resource, encodedParams.toArray());
        }
        return path;
    }

    private ListenableFuture<Response> request(final URI uri, final String method) {
        return request(uri, method, null);
    }

    private ListenableFuture<Response> request(final URI uri, final String method, final Object entity) {
        final Map<String, List<String>> headers = Maps.newHashMap();
        final byte[] entityBytes;
        headers.put(VersionCompatibility.HELIOS_VERSION_HEADER, Collections.singletonList(Version.POM_VERSION));
        if (entity != null) {
            headers.put("Content-Type", singletonList("application/json"));
            headers.put("Charset", singletonList("utf-8"));
            entityBytes = Json.asBytesUnchecked(entity);
        } else {
            entityBytes = new byte[] {};
        }

        final ListenableFuture<Response> f = dispatcher.request(uri, method, entityBytes, headers);
        return transform(f, new Function<Response, Response>() {
            @Override
            public Response apply(final Response response) {
                checkProtocolVersionStatus(response);
                return response;
            }
        });
    }

    private void checkProtocolVersionStatus(final Response response) {
        final VersionCompatibility.Status versionStatus = getVersionStatus(response);
        if (versionStatus == null) {
            log.debug("Server didn't return a version header!");
            return; // shouldn't happen really
        }

        final String serverVersion = response.header(HELIOS_SERVER_VERSION_HEADER);
        if ((versionStatus == VersionCompatibility.Status.MAYBE)
                && (versionWarningLogged.compareAndSet(false, true))) {
            log.warn("Your Helios client version [{}] is ahead of the server [{}].  This will"
                    + " probably work ok but there is the potential for weird things.  If in doubt,"
                    + " contact the Helios team if you think the cluster you're connecting to is out"
                    + " of date and should be upgraded.", Version.POM_VERSION, serverVersion);
        }
    }

    private VersionCompatibility.Status getVersionStatus(final Response response) {
        final String status = response.header(HELIOS_VERSION_STATUS_HEADER);
        if (status != null) {
            return VersionCompatibility.Status.valueOf(status);
        }
        return null;
    }

    private <T> ListenableFuture<T> get(final URI uri, final TypeReference<T> typeReference) {
        return get(uri, Json.type(typeReference));
    }

    private <T> ListenableFuture<T> get(final URI uri, final Class<T> clazz) {
        return get(uri, Json.type(clazz));
    }

    private <T> ListenableFuture<T> get(final URI uri, final JavaType javaType) {
        return transform(request(uri, "GET"), new ConvertResponseToPojo<T>(javaType));
    }

    private ListenableFuture<Integer> put(final URI uri) {
        return status(request(uri, "PUT"));
    }

    public ListenableFuture<JobDeployResponse> deploy(final Deployment job, final String host) {
        return deploy(job, host, "");
    }

    public ListenableFuture<JobDeployResponse> deploy(final Deployment job, final String host, final String token) {
        final Set<Integer> deserializeReturnCodes = ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND, HTTP_BAD_METHOD,
                HTTP_BAD_REQUEST, HTTP_FORBIDDEN);
        return transform(
                request(uri(path("/hosts/%s/jobs/%s", host, job.getJobId()), ImmutableMap.of("token", token)),
                        "PUT", job),
                ConvertResponseToPojo.create(JobDeployResponse.class, deserializeReturnCodes));
    }

    public ListenableFuture<SetGoalResponse> setGoal(final Deployment job, final String host) {
        return setGoal(job, host, "");
    }

    public ListenableFuture<SetGoalResponse> setGoal(final Deployment job, final String host, final String token) {
        return transform(
                request(uri(path("/hosts/%s/jobs/%s", host, job.getJobId()), ImmutableMap.of("token", token)),
                        "PATCH", job),
                ConvertResponseToPojo.create(SetGoalResponse.class,
                        ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND, HTTP_FORBIDDEN)));
    }

    private ListenableFuture<Integer> status(final ListenableFuture<Response> req) {
        return transform(req, new Function<Response, Integer>() {
            @Override
            public Integer apply(final Response reply) {
                return reply.status();
            }
        });
    }

    public ListenableFuture<Deployment> deployment(final String host, final JobId job) {
        return get(uri(path("/hosts/%s/jobs/%s", host, job)), Deployment.class);
    }

    public ListenableFuture<HostStatus> hostStatus(final String host) {
        return hostStatus(host, Collections.<String, String>emptyMap());
    }

    public ListenableFuture<HostStatus> hostStatus(final String host, final Map<String, String> queryParams) {
        return get(uri(path("/hosts/%s/status", host), queryParams), HostStatus.class);
    }

    public ListenableFuture<Map<String, HostStatus>> hostStatuses(final List<String> hosts) {
        return hostStatuses(hosts, Collections.<String, String>emptyMap());
    }

    public ListenableFuture<Map<String, HostStatus>> hostStatuses(final List<String> hosts,
            final Map<String, String> queryParams) {
        final ConvertResponseToPojo<Map<String, HostStatus>> converter = ConvertResponseToPojo.create(
                TypeFactory.defaultInstance().constructMapType(Map.class, String.class, HostStatus.class),
                ImmutableSet.of(HTTP_OK));

        return transform(request(uri("/hosts/statuses", queryParams), "POST", hosts), converter);
    }

    public ListenableFuture<Integer> registerHost(final String host, final String id) {
        return put(uri(path("/hosts/%s", host), ImmutableMap.of("id", id)));
    }

    public ListenableFuture<JobDeleteResponse> deleteJob(final JobId id) {
        return deleteJob(id, "");
    }

    public ListenableFuture<JobDeleteResponse> deleteJob(final JobId id, final String token) {
        return transform(request(uri(path("/jobs/%s", id), ImmutableMap.of("token", token)), "DELETE"),
                ConvertResponseToPojo.create(JobDeleteResponse.class,
                        ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND, HTTP_BAD_REQUEST, HTTP_FORBIDDEN)));
    }

    public ListenableFuture<JobUndeployResponse> undeploy(final JobId jobId, final String host) {
        return undeploy(jobId, host, "");
    }

    public ListenableFuture<JobUndeployResponse> undeploy(final JobId jobId, final String host,
            final String token) {
        return transform(
                request(uri(path("/hosts/%s/jobs/%s", host, jobId), ImmutableMap.of("token", token)), "DELETE"),
                ConvertResponseToPojo.create(JobUndeployResponse.class,
                        ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND, HTTP_BAD_REQUEST, HTTP_FORBIDDEN)));
    }

    public ListenableFuture<HostDeregisterResponse> deregisterHost(final String host) {
        return transform(request(uri(path("/hosts/%s", host)), "DELETE"), ConvertResponseToPojo
                .create(HostDeregisterResponse.class, ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND)));
    }

    public ListenableFuture<List<String>> listHosts() {
        return get(uri("/hosts/"), new TypeReference<List<String>>() {
        });
    }

    public ListenableFuture<List<String>> listMasters() {
        return get(uri("/masters/"), new TypeReference<List<String>>() {
        });
    }

    public ListenableFuture<VersionResponse> version() {
        // Create a fallback in case we fail to connect to the master. Return null if this happens.
        // The transform below will handle this and return an appropriate error message to the caller.
        final ListenableFuture<Response> futureWithFallback = withFallback(request(uri("/version/"), "GET"),
                new FutureFallback<Response>() {
                    @Override
                    public ListenableFuture<Response> create(@NotNull Throwable t) throws Exception {
                        return immediateFuture(null);
                    }
                });

        return transform(futureWithFallback, new AsyncFunction<Response, VersionResponse>() {
            @Override
            public ListenableFuture<VersionResponse> apply(@NotNull Response reply) throws Exception {
                final String masterVersion = reply == null ? "Unable to connect to master"
                        : reply.status() == HTTP_OK ? Json.read(reply.payload(), String.class)
                                : "Master replied with error code " + reply.status();

                return immediateFuture(new VersionResponse(Version.POM_VERSION, masterVersion));
            }
        });
    }

    public ListenableFuture<CreateJobResponse> createJob(final Job descriptor) {
        return transform(request(uri("/jobs/"), "POST", descriptor),
                ConvertResponseToPojo.create(CreateJobResponse.class, ImmutableSet.of(HTTP_OK, HTTP_BAD_REQUEST)));
    }

    public ListenableFuture<Map<JobId, Job>> jobs(final String query) {
        return get(uri("/jobs", ImmutableMap.of("q", query)), new TypeReference<Map<JobId, Job>>() {
        });
    }

    public ListenableFuture<Map<JobId, Job>> jobs() {
        return get(uri("/jobs"), new TypeReference<Map<JobId, Job>>() {
        });
    }

    public ListenableFuture<TaskStatusEvents> jobHistory(final JobId jobId) {
        return transform(request(uri(path("/history/jobs/%s", jobId.toString())), "GET"),
                ConvertResponseToPojo.create(TaskStatusEvents.class, ImmutableSet.of(HTTP_OK, HTTP_NOT_FOUND)));
    }

    public ListenableFuture<JobStatus> jobStatus(final JobId jobId) {
        return get(uri(path("/jobs/%s/status", jobId)), JobStatus.class);
    }

    public ListenableFuture<Map<JobId, JobStatus>> jobStatuses(final Set<JobId> jobs) {
        final ConvertResponseToPojo<Map<JobId, JobStatus>> converter = ConvertResponseToPojo.create(
                TypeFactory.defaultInstance().constructMapType(Map.class, JobId.class, JobStatus.class),
                ImmutableSet.of(HTTP_OK));

        return transform(request(uri("/jobs/statuses"), "POST", jobs), converter);
    }

    public ListenableFuture<DeploymentGroup> deploymentGroup(final String name) {
        return get(uri("/deployment-group/" + name), new TypeReference<DeploymentGroup>() {
        });
    }

    public ListenableFuture<List<String>> listDeploymentGroups() {
        return get(uri("/deployment-group/"), new TypeReference<List<String>>() {
        });
    }

    public ListenableFuture<DeploymentGroupStatusResponse> deploymentGroupStatus(final String name) {
        return get(uri(path("/deployment-group/%s/status", name)),
                new TypeReference<DeploymentGroupStatusResponse>() {
                });
    }

    public ListenableFuture<CreateDeploymentGroupResponse> createDeploymentGroup(final DeploymentGroup descriptor) {
        return transform(request(uri("/deployment-group/"), "POST", descriptor), ConvertResponseToPojo
                .create(CreateDeploymentGroupResponse.class, ImmutableSet.of(HTTP_OK, HTTP_BAD_REQUEST)));
    }

    public ListenableFuture<RemoveDeploymentGroupResponse> removeDeploymentGroup(final String name) {
        return transform(request(uri("/deployment-group/" + name), "DELETE"), ConvertResponseToPojo
                .create(RemoveDeploymentGroupResponse.class, ImmutableSet.of(HTTP_OK, HTTP_BAD_REQUEST)));
    }

    public ListenableFuture<RollingUpdateResponse> rollingUpdate(final String deploymentGroupName, final JobId job,
            final RolloutOptions options) {
        return transform(
                request(uri(path("/deployment-group/%s/rolling-update", deploymentGroupName)), "POST",
                        new RollingUpdateRequest(job, options)),
                ConvertResponseToPojo.create(RollingUpdateResponse.class,
                        ImmutableSet.of(HTTP_OK, HTTP_BAD_REQUEST)));
    }

    public ListenableFuture<Integer> stopDeploymentGroup(final String deploymentGroupName) {
        return status(request(uri(path("/deployment-group/%s/stop", deploymentGroupName)), "POST"));
    }

    private static final class ConvertResponseToPojo<T> implements AsyncFunction<Response, T> {

        private final JavaType javaType;
        private final Set<Integer> decodeableStatusCodes;

        private ConvertResponseToPojo(final JavaType javaType) {
            this(javaType, ImmutableSet.of(HTTP_OK));
        }

        public ConvertResponseToPojo(final JavaType type, final Set<Integer> decodeableStatusCodes) {
            this.javaType = type;
            this.decodeableStatusCodes = decodeableStatusCodes;
        }

        public static <T> ConvertResponseToPojo<T> create(final JavaType type,
                final Set<Integer> decodeableStatusCodes) {
            return new ConvertResponseToPojo<>(type, decodeableStatusCodes);
        }

        public static <T> ConvertResponseToPojo<T> create(final Class<T> clazz,
                final Set<Integer> decodeableStatusCodes) {
            return new ConvertResponseToPojo<>(Json.type(clazz), decodeableStatusCodes);
        }

        @Override
        public ListenableFuture<T> apply(@NotNull final Response reply) throws HeliosException {
            if (reply.status() == HTTP_NOT_FOUND && !decodeableStatusCodes.contains(HTTP_NOT_FOUND)) {
                return immediateFuture(null);
            }

            if (!decodeableStatusCodes.contains(reply.status())) {
                throw new HeliosException("request failed: " + reply);
            }

            if (reply.payload().length == 0) {
                throw new HeliosException("bad reply: " + reply);
            }

            final T result;
            try {
                result = Json.read(reply.payload(), javaType);
            } catch (IOException e) {
                throw new HeliosException("bad reply: " + reply, e);
            }

            return immediateFuture(result);
        }
    }

    public static Builder newBuilder() {
        return new Builder();
    }

    public static class Builder {

        private static final String HELIOS_CERT_PATH = "HELIOS_CERT_PATH";

        private String user;
        private ClientCertificatePath clientCertificatePath;
        private Supplier<List<Endpoint>> endpointSupplier;
        private boolean sslHostnameVerification = true;

        private Builder() {
        }

        public Builder setUser(final String user) {
            this.user = user;
            return this;
        }

        public Builder setDomain(final String domain) {
            return setEndpointSupplier(Endpoints.of(Resolver.supplier("helios", domain)));
        }

        public Builder setEndpoints(final List<URI> endpoints) {
            return setEndpointSupplier(Suppliers.ofInstance(Endpoints.of(endpoints)));
        }

        public Builder setEndpoints(final URI... endpoints) {
            return setEndpointSupplier(Suppliers.ofInstance(Endpoints.of(asList(endpoints))));
        }

        public Builder setEndpoints(final String... endpoints) {
            return setEndpointStrings(asList(endpoints));
        }

        public Builder setEndpointStrings(final List<String> endpoints) {
            final List<URI> uris = Lists.newArrayList();
            for (String endpoint : endpoints) {
                uris.add(URI.create(endpoint));
            }
            return setEndpoints(uris);
        }

        public Builder setEndpointSupplier(final Supplier<List<Endpoint>> endpointSupplier) {
            this.endpointSupplier = endpointSupplier;
            return this;
        }

        /**
         * Can be used to disable hostname verification for HTTPS connections to the Helios master.
         * Defaults to being enabled.
         */
        public Builder setSslHostnameVerification(final boolean enabled) {
            this.sslHostnameVerification = enabled;
            return this;
        }

        public Builder setClientCertificatePath(final ClientCertificatePath clientCertificatePath) {
            this.clientCertificatePath = clientCertificatePath;
            return this;
        }

        public HeliosClient build() {
            return new HeliosClient(user, createDispatcher());
        }

        private RequestDispatcher createDispatcher() {
            final ScheduledExecutorService executor = MoreExecutors.getExitingScheduledExecutorService(
                    (ScheduledThreadPoolExecutor) newScheduledThreadPool(4), 0, SECONDS);

            final ListeningScheduledExecutorService listeningExecutor = MoreExecutors.listeningDecorator(executor);

            final RequestDispatcher dispatcher = new DefaultRequestDispatcher(
                    createHttpConnector(sslHostnameVerification), listeningExecutor);

            return new RetryingRequestDispatcher(dispatcher, listeningExecutor, new SystemClock(), 5,
                    TimeUnit.SECONDS);
        }

        private HttpConnector createHttpConnector(final boolean sslHostnameVerification) {

            final EndpointIterator endpointIterator = EndpointIterator.of(endpointSupplier.get());

            final DefaultHttpConnector connector = new DefaultHttpConnector(endpointIterator, 10000,
                    sslHostnameVerification);

            Optional<AgentProxy> agentProxyOpt = Optional.absent();
            try {
                agentProxyOpt = Optional.of(AgentProxies.newInstance());
            } catch (RuntimeException e) {
                // the user likely doesn't have ssh-agent setup. This may not matter at all if the masters
                // do not require authentication, so we delay reporting any sort of error to the user until
                // the servers return 401 Unauthorized.
                log.debug("{}", e);
            }

            // set up the ClientCertificatePath, giving precedence to any values set
            // with setClientCertificatePath()
            ClientCertificatePath clientCertificatePath = this.clientCertificatePath;
            if (clientCertificatePath == null) {
                final String heliosCertPath = System.getenv(HELIOS_CERT_PATH);
                if (!isNullOrEmpty(heliosCertPath)) {
                    final Path certPath = Paths.get(heliosCertPath, "cert.pem");
                    final Path keyPath = Paths.get(heliosCertPath, "key.pem");

                    if (certPath.toFile().canRead() && keyPath.toFile().canRead()) {
                        this.clientCertificatePath = new ClientCertificatePath(certPath, keyPath);
                    } else {
                        log.warn(
                                "{} is set to {}, but {} and/or {} do not exist or cannot be read. "
                                        + "Will not send client certificate in HeliosClient requests.",
                                HELIOS_CERT_PATH, heliosCertPath, certPath, keyPath);
                    }
                }
            }

            return new AuthenticatingHttpConnector(user, agentProxyOpt,
                    Optional.fromNullable(clientCertificatePath), endpointIterator, connector);
        }
    }

    /**
     * Create a new helios client as a specific user, connecting to a helios master cluster in a
     * specific domain.
     *
     * @param domain The target domain.
     * @param user   The user to identify as.
     * @return A helios client.
     */
    public static HeliosClient create(final String domain, final String user) {
        return HeliosClient.newBuilder().setDomain(domain).setUser(user).build();
    }

}