LauncherLifecycleManager.java

package io.airlift.airship.agent;

import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.Files;
import com.google.common.net.InetAddresses;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.Inject;
import io.airlift.airship.shared.SlotLifecycleState;
import io.airlift.command.Command;
import io.airlift.command.CommandFailedException;
import io.airlift.http.server.HttpServerInfo;
import io.airlift.log.Logger;
import io.airlift.node.NodeInfo;
import io.airlift.units.Duration;

import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.net.URI;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

import static io.airlift.airship.shared.HttpUriBuilder.uriBuilderFrom;
import static io.airlift.airship.shared.SlotLifecycleState.RUNNING;
import static io.airlift.airship.shared.SlotLifecycleState.STOPPED;
import static io.airlift.airship.shared.SlotLifecycleState.UNKNOWN;

public class LauncherLifecycleManager
        implements LifecycleManager
{
    private static final Logger log = Logger.get(LauncherLifecycleManager.class);

    private final Executor executor;
    private final InetAddress internalIp;
    private final String externalAddress;
    private final Duration launcherTimeout;
    private final Duration stopTimeout;
    private final String environment;
    private final InetAddress bindIp;
    private final URI serviceInventoryUri;

    @Inject
    public LauncherLifecycleManager(AgentConfig config, NodeInfo nodeInfo, HttpServerInfo httpServerInfo)
    {
        this(nodeInfo.getEnvironment(),
                nodeInfo.getInternalIp(),
                nodeInfo.getExternalAddress(),
                nodeInfo.getBindIp(),
                config.getLauncherTimeout(),
                config.getLauncherStopTimeout(),
                uriBuilderFrom(httpServerInfo.getHttpsUri() != null ? httpServerInfo.getHttpsUri() : httpServerInfo.getHttpUri()).appendPath("/v1/serviceInventory").build()
        );
    }

    public LauncherLifecycleManager(String environment,
            InetAddress internalIp,
            String externalAddress,
            InetAddress bindIp,
            Duration launcherTimeout,
            Duration launcherStopTimeout,
            URI serviceInventoryUri)
    {
        this.launcherTimeout = launcherTimeout;
        stopTimeout = launcherStopTimeout;

        executor = Executors.newCachedThreadPool(new ThreadFactoryBuilder().setDaemon(true).setNameFormat("launcher-command-%s").build());
        this.environment = environment;
        this.internalIp = internalIp;
        this.externalAddress = externalAddress;
        this.bindIp = bindIp;

        this.serviceInventoryUri = serviceInventoryUri;
    }

    @Override
    public SlotLifecycleState status(Deployment deployment)
    {
        try {
            int exitCode = createCommand("status", deployment, launcherTimeout)
                    .setSuccessfulExitCodes(0, 1, 2, 3)
                    .execute(executor);
            if (exitCode == 0) {
                return RUNNING;
            }
            else {
                return STOPPED;
            }
        }
        catch (CommandFailedException e) {
            return UNKNOWN;
        }
    }

    @Override
    public SlotLifecycleState start(Deployment deployment)
    {
        updateNodeConfig(deployment);
        Command command = createCommand("start", deployment, launcherTimeout);
        try {
            command.execute(executor);
            return RUNNING;
        }
        catch (CommandFailedException e) {
            log.error("ENVIRONMENT:\n    %s", Joiner.on("\n    ").withKeyValueSeparator("=").join(command.getEnvironment()));
            throw new RuntimeException("start failed: " + e.getMessage());
        }
    }

    @Override
    public SlotLifecycleState restart(Deployment deployment)
    {
        updateNodeConfig(deployment);
        try {
            Command command = createCommand("restart", deployment, stopTimeout);
            command.execute(executor);
            return RUNNING;
        }
        catch (CommandFailedException e) {
            throw new RuntimeException("restart failed: " + e.getMessage());
        }
    }

    @Override
    public SlotLifecycleState stop(Deployment deployment)
    {
        updateNodeConfig(deployment);
        try {
            createCommand("stop", deployment, stopTimeout).execute(executor);
            return STOPPED;
        }
        catch (CommandFailedException e) {
            throw new RuntimeException("stop failed: " + e.getMessage());
        }
    }

    @Override
    public SlotLifecycleState kill(Deployment deployment)
    {
        updateNodeConfig(deployment);
        try {
            createCommand("kill", deployment, launcherTimeout).execute(executor);
            return STOPPED;
        }
        catch (CommandFailedException e) {
            throw new RuntimeException("kill failed: " + e.getMessage());
        }
    }

    private Command createCommand(String commandName, Deployment deployment, Duration timeLimit)
    {

        File launcherScript = new File(new File(deployment.getDeploymentDir(), "bin"), "launcher");

        Command command = new Command(launcherScript.getAbsolutePath(), commandName)
                .setDirectory(deployment.getDataDir())
                .setTimeLimit(timeLimit)
                .addEnvironment("HOME", deployment.getDataDir().getAbsolutePath());

        return command;
    }

    @Override
    public void updateNodeConfig(Deployment deployment)
    {
        ImmutableMap.Builder<String, String> map = ImmutableMap.builder();

        map.put("node.environment", environment);
        map.put("node.id", deployment.getNodeId().toString());
        map.put("node.location", deployment.getLocation());
        map.put("node.data-dir", deployment.getDataDir().getAbsolutePath());
        map.put("node.binary-spec", deployment.getAssignment().getBinary());
        map.put("node.config-spec", deployment.getAssignment().getConfig());

        if (internalIp != null) {
            map.put("node.ip", InetAddresses.toAddrString(internalIp));
        }

        if (externalAddress != null) {
            map.put("node.external-address", externalAddress);
        }

        // add bind ip only if explicitly set on the agent
        if (bindIp != null && InetAddresses.coerceToInteger(bindIp) != 0) {
            map.put("node.bind-ip", bindIp.getHostAddress());
        }

        // add service inventory uri
        map.put("service-inventory.uri", serviceInventoryUri.toString());

        File nodeConfig = new File(deployment.getDeploymentDir(), "etc/node.properties");
        nodeConfig.getParentFile().mkdirs();

        try {
            String data = Joiner.on("\n").withKeyValueSeparator("=").join(map.build()) + "\n";
            Files.write(data, nodeConfig, Charsets.UTF_8);
        }
        catch (IOException e) {
            nodeConfig.delete();
            throw new RuntimeException("create node config failed: " + e.getMessage());
        }
    }
}