/*
* Copyright (c) 2014 Spotify AB.
*
* 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 com.spotify.helios.testing;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.common.io.Files;
import com.google.common.io.Resources;
import com.fasterxml.jackson.databind.JsonNode;
import com.spotify.helios.common.Json;
import com.spotify.helios.common.descriptors.Job;
import com.spotify.helios.common.descriptors.PortMapping;
import com.spotify.helios.common.descriptors.ServiceEndpoint;
import com.spotify.helios.common.descriptors.ServicePorts;
import org.apache.commons.lang.text.StrSubstitutor;
import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.file.Path;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.regex.Pattern;
import static com.fasterxml.jackson.databind.node.JsonNodeType.STRING;
import static com.google.common.base.Charsets.UTF_8;
import static com.google.common.base.Optional.fromNullable;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.isNullOrEmpty;
import static java.lang.Integer.toHexString;
import static java.lang.System.getenv;
import static java.util.Arrays.asList;
import static org.junit.Assert.fail;
public class TemporaryJobBuilder {
private static final Logger log = LoggerFactory.getLogger(TemporaryJobBuilder.class);
private static final Pattern JOB_NAME_FORBIDDEN_CHARS = Pattern.compile("[^0-9a-zA-Z-_.]+");
private static final int DEFAULT_EXPIRES_MINUTES = 30;
private final List<String> hosts = Lists.newArrayList();
private final Job.Builder builder = Job.newBuilder();
private final Set<String> waitPorts = Sets.newHashSet();
private final Deployer deployer;
private final String jobNamePrefix;
private final String jobDeployedMessageFormat;
private String hostFilter;
private Prober prober;
private TemporaryJob job;
public TemporaryJobBuilder(final Deployer deployer, final String jobNamePrefix,
final Prober defaultProber) {
this(deployer, jobNamePrefix, defaultProber, (String) null);
}
public TemporaryJobBuilder(final Deployer deployer, final String jobNamePrefix,
final Prober defaultProber, final String jobDeployedMessageFormat) {
checkNotNull(deployer, "deployer");
checkNotNull(jobNamePrefix, "jobNamePrefix");
checkNotNull(defaultProber, "defaultProber");
this.deployer = deployer;
this.jobNamePrefix = jobNamePrefix;
this.prober = defaultProber;
this.builder.setRegistrationDomain(jobNamePrefix);
this.jobDeployedMessageFormat = jobDeployedMessageFormat;
}
public TemporaryJobBuilder name(final String jobName) {
this.builder.setName(jobName);
return this;
}
public TemporaryJobBuilder version(final String jobVersion) {
this.builder.setVersion(jobVersion);
return this;
}
public TemporaryJobBuilder image(final String image) {
this.builder.setImage(image);
return this;
}
public TemporaryJobBuilder registrationDomain(final String domain) {
this.builder.setRegistrationDomain(domain);
return this;
}
public TemporaryJobBuilder command(final List<String> command) {
this.builder.setCommand(command);
return this;
}
public TemporaryJobBuilder command(final String... command) {
return command(asList(command));
}
public TemporaryJobBuilder env(final String key, final Object value) {
this.builder.addEnv(key, value.toString());
return this;
}
public TemporaryJobBuilder disablePrivateRegistrationDomain() {
this.builder.setRegistrationDomain(Job.EMPTY_REGISTRATION_DOMAIN);
return this;
}
public TemporaryJobBuilder port(final String name, final int internalPort) {
return port(name, internalPort, true);
}
public TemporaryJobBuilder port(final String name, final int internalPort, final boolean wait) {
return port(name, internalPort, null, wait);
}
public TemporaryJobBuilder port(final String name, final int internalPort,
final Integer externalPort) {
return port(name, internalPort, externalPort, true);
}
public TemporaryJobBuilder port(final String name, final int internalPort,
final Integer externalPort, final boolean wait) {
this.builder.addPort(name, PortMapping.of(internalPort, externalPort));
if (wait) {
waitPorts.add(name);
}
return this;
}
public TemporaryJobBuilder registration(final ServiceEndpoint endpoint,
final ServicePorts ports) {
this.builder.addRegistration(endpoint, ports);
return this;
}
public TemporaryJobBuilder registration(final String service, final String protocol,
final String... ports) {
return registration(ServiceEndpoint.of(service, protocol), ServicePorts.of(ports));
}
public TemporaryJobBuilder registration(final Map<ServiceEndpoint, ServicePorts> registration) {
this.builder.setRegistration(registration);
return this;
}
public TemporaryJobBuilder host(final String host) {
this.hosts.add(host);
return this;
}
public TemporaryJobBuilder hostFilter(final String hostFilter) {
this.hostFilter = hostFilter;
return this;
}
/**
* The Helios master will undeploy and delete the job at the specified date, if it has not
* already been removed. If not set, jobs will be removed after 30 minutes. This is for the
* case when a TemporaryJob is not cleaned up properly, perhaps because the process terminated
* prematurely.
* @param expires the Date when the job should be removed
* @return the TemporaryJobBuilder
*/
public TemporaryJobBuilder expires(final Date expires) {
this.builder.setExpires(expires);
return this;
}
/**
* This will override the default prober provided by {@link TemporaryJobs} to the constructor.
* @param prober the prober to use for this job
* @return the TemporaryJobBuilder
*/
public TemporaryJobBuilder prober(final Prober prober) {
this.prober = prober;
return this;
}
/**
* Deploys the job to the specified hosts. If no hosts are specified, a host will be chosen at
* random from the current Helios cluster. If the HELIOS_HOST_FILTER environment variable is set,
* it will be used to filter the list of hosts in the current Helios cluster.
*
* @param hosts the list of helios hosts to deploy to. A random host will be chosen if the list is
* empty.
* @return a TemporaryJob representing the deployed job
*/
public TemporaryJob deploy(final String... hosts) {
return deploy(asList(hosts));
}
/**
* Deploys the job to the specified hosts. If no hosts are specified, a host will be chosen at
* random from the current Helios cluster. If the HELIOS_HOST_FILTER environment variable is set,
* it will be used to filter the list of hosts in the current Helios cluster.
*
* @param hosts the list of helios hosts to deploy to. A random host will be chosen if the list is
* empty.
* @return a TemporaryJob representing the deployed job
*/
public TemporaryJob deploy(final List<String> hosts) {
this.hosts.addAll(hosts);
if (job == null) {
if (builder.getName() == null && builder.getVersion() == null) {
// Both name and version are unset, use image name as job name and generate random version
builder.setName(jobName(builder.getImage(), jobNamePrefix));
builder.setVersion(randomVersion());
}
// Set job to expires value, if it's not already set. This ensures temporary jobs which
// aren't cleaned up properly by the test will be removed by the master.
if (builder.getExpires() == null) {
builder.setExpires(new DateTime().plusMinutes(DEFAULT_EXPIRES_MINUTES).toDate());
}
if (this.hosts.isEmpty()) {
if (isNullOrEmpty(hostFilter)) {
hostFilter = getenv("HELIOS_HOST_FILTER");
}
job = deployer.deploy(builder.build(), hostFilter, waitPorts, prober);
} else {
job = deployer.deploy(builder.build(), this.hosts, waitPorts, prober);
}
if (!Strings.isNullOrEmpty(jobDeployedMessageFormat)) {
outputMessage(job);
}
}
return job;
}
private void outputMessage(final TemporaryJob job) {
for (String host : job.hosts()) {
final StrSubstitutor subst = new StrSubstitutor(new ImmutableMap.Builder<String, Object>()
.put("host", host)
.put("name", job.job().getId().getName())
.put("version", job.job().getId().getVersion())
.put("hash", job.job().getId().getHash())
.put("job", job.job().toString())
.put("containerId", job.statuses().get(host).getContainerId())
.build()
);
log.info("{}", subst.replace(jobDeployedMessageFormat));
}
}
public TemporaryJobBuilder imageFromBuild() {
final String envPath = getenv("IMAGE_INFO_PATH");
if (envPath != null) {
return imageFromInfoFile(envPath);
} else {
final String name = fromNullable(getenv("IMAGE_INFO_NAME")).or("image_info.json");
URL info;
try {
info = Resources.getResource(name);
} catch (IllegalArgumentException e) {
info = getFromFileSystem(name);
if (info == null) {
throw e;
}
}
try {
final String json = Resources.asCharSource(info, UTF_8).read();
return imageFromInfoJson(json, info.toString());
} catch (IOException e) {
throw new AssertionError("Failed to load image info", e);
}
}
}
private URL getFromFileSystem(String name) {
final File file = new File("target/" + name);
if (!file.exists()) {
return null;
}
try {
return file.toURI().toURL();
} catch (MalformedURLException e) {
throw new RuntimeException("TEST");
}
}
public TemporaryJobBuilder imageFromInfoFile(final Path path) {
return imageFromInfoFile(path.toFile());
}
public TemporaryJobBuilder imageFromInfoFile(final String path) {
return imageFromInfoFile(new File(path));
}
public TemporaryJobBuilder imageFromInfoFile(final File file) {
final String json;
try {
json = Files.toString(file, UTF_8);
} catch (IOException e) {
throw new AssertionError("Failed to read image info file: " +
file + ": " + e.getMessage());
}
return imageFromInfoJson(json, file.toString());
}
private TemporaryJobBuilder imageFromInfoJson(final String json,
final String source) {
try {
final JsonNode info = Json.readTree(json);
final JsonNode imageNode = info.get("image");
if (imageNode == null) {
fail("Missing image field in image info: " + source);
}
if (imageNode.getNodeType() != STRING) {
fail("Bad image field in image info: " + source);
}
final String image = imageNode.asText();
return image(image);
} catch (IOException e) {
throw new AssertionError("Failed to parse image info: " + source, e);
}
}
private String jobName(final String s, final String jobNamePrefix) {
return jobNamePrefix + "_" + JOB_NAME_FORBIDDEN_CHARS.matcher(s).replaceAll("_");
}
private String randomVersion() {
return toHexString(ThreadLocalRandom.current().nextInt());
}
}