Merge pull request #4707 from vespa-engine/hmusum/file-distribution-status-client

Add client for file distribution API

3 files changed, 335 insertions, 0 deletions

diff --git a/filedistribution/pom.xml b/filedistribution/pom.xml
index d9699b700d0..6bfa4362f65 100644
--- a/filedistribution/pom.xml
+++ b/filedistribution/pom.xml
@@ -58,6 +58,18 @@
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-compress</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>airline</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.httpcomponents</groupId>
+            <artifactId>httpclient</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+        </dependency>
     </dependencies>
 
     <build>
@@ -67,6 +79,34 @@
                 <artifactId>bundle-plugin</artifactId>
                 <extensions>true</extensions>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <configuration>
+                    <compilerArgs>
+                        <arg>-Xlint:all</arg>
+                        <arg>-Xlint:-serial</arg>
+                        <arg>-Werror</arg>
+                    </compilerArgs>
+                </configuration>
+            </plugin>
+            <plugin>
+                <artifactId>maven-assembly-plugin</artifactId>
+                <configuration>
+                    <descriptorRefs>
+                        <descriptorRef>jar-with-dependencies</descriptorRef>
+                    </descriptorRefs>
+                </configuration>
+                <executions>
+                    <execution>
+                        <id>make-assembly</id>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>single</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
         </plugins>
     </build>
 
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClient.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClient.java
new file mode 100644
index 00000000000..b50416ac159
--- /dev/null
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClient.java
@@ -0,0 +1,236 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.filedistribution.status;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.airlift.airline.Command;
+import io.airlift.airline.HelpOption;
+import io.airlift.airline.Option;
+import io.airlift.airline.SingleCommand;
+import org.apache.http.client.config.RequestConfig;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.client.utils.URIBuilder;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.util.EntityUtils;
+
+import javax.inject.Inject;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.Map;
+
+/**
+ * Tool for getting file distribution status
+ *
+ * @author hmusum
+ */
+public class FileDistributionStatusClient {
+
+    private static final String statusUnknown = "UNKNOWN";
+    private static final String statusInProgress = "IN_PROGRESS";
+    private static final String statusFinished = "FINISHED";
+
+
+    private final String tenantName;
+    private final String applicationName;
+    private final String instanceName;
+    private final String environment;
+    private final String region;
+    private final double timeout;
+    private final boolean debug;
+
+    FileDistributionStatusClient(CommandLineArguments arguments) {
+        tenantName = arguments.getTenantName();
+        applicationName = arguments.getApplicationName();
+        instanceName = arguments.getInstanceName();
+        environment = arguments.getEnvironment();
+        region = arguments.getRegion();
+        timeout = arguments.getTimeout();
+        debug = arguments.getDebugFlag();
+    }
+
+    public static void main(String[] args) {
+        try {
+            new FileDistributionStatusClient(CommandLineArguments.build(args)).run();
+        } catch (Exception e) {
+            System.err.println(e.getMessage());
+            System.exit(1);
+        }
+    }
+
+    public void run() {
+        String json = doHttpRequest();
+        System.out.println(parseAndGenerateOutput(json));
+    }
+
+    private String doHttpRequest() {
+        int timeoutInMillis = (int) (timeout * 1000);
+        RequestConfig config = RequestConfig.custom()
+                .setConnectTimeout(timeoutInMillis)
+                .setConnectionRequestTimeout(timeoutInMillis)
+                .setSocketTimeout(timeoutInMillis)
+                .build();
+        CloseableHttpClient httpClient = HttpClientBuilder.create().setDefaultRequestConfig(config).build();
+        URI statusUri = createStatusApiUri();
+        if (debug)
+            System.out.println("URI:" + statusUri);
+        try {
+            CloseableHttpResponse response = httpClient.execute(new HttpGet(statusUri));
+            String content = EntityUtils.toString(response.getEntity());
+            if (debug)
+                System.out.println("response:" + content);
+            if (response.getStatusLine().getStatusCode() == 200) {
+                return content;
+            } else {
+                throw new RuntimeException("Failed to get status for request " + statusUri + ": " +
+                                                   response.getStatusLine() + ": " + content);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    String parseAndGenerateOutput(String json) {
+        ObjectMapper objectMapper = new ObjectMapper();
+        JsonNode jsonNode;
+        try {
+            jsonNode = objectMapper.readTree(json);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        String status = jsonNode.get("status").asText();
+        switch (status) {
+            case statusUnknown:
+                return "File distribution status unknown: " + jsonNode.get("message").asText();
+            case statusInProgress:
+                return "File distribution in progress:\n" + inProgressOutput(jsonNode.get("hosts"));
+            case statusFinished:
+                return "File distribution finished";
+            default:
+                throw new RuntimeException("Unknown status " + status);
+        }
+    }
+
+    private URI createStatusApiUri() {
+        String path = String.format("/application/v2/tenant/%s/application/%s/environment/%s/region/%s/instance/%s/filedistributionstatus",
+                                    tenantName, applicationName, environment, region, instanceName);
+        try {
+            return new URIBuilder()
+                    .setScheme("http")
+                    .setHost("localhost")
+                    .setPort(19071)
+                    .setPath(path)
+                    .addParameter("timeout", String.valueOf(timeout))
+                    .build();
+        } catch (URISyntaxException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private String inProgressOutput(JsonNode hosts) {
+        ArrayList<String> statusPerHost = new ArrayList<>();
+        for (JsonNode host : hosts) {
+            StringBuilder sb = new StringBuilder();
+            String status = host.get("status").asText();
+            sb.append(host.get("hostname").asText()).append(": ").append(status);
+            if (status.equals(statusUnknown))
+                sb.append(" (").append(host.get("message").asText()).append(")");
+            else if (status.equals(statusInProgress)) {
+                JsonNode fileReferencesArray = host.get("fileReferences");
+                int size = fileReferencesArray.size();
+                int finished = 0;
+                for (JsonNode element : fileReferencesArray) {
+                    for (Iterator<Map.Entry<String, JsonNode>> it = element.fields(); it.hasNext(); ) {
+                        Map.Entry<String, JsonNode> fileReferenceStatus = it.next();
+                        if (fileReferenceStatus.getValue().asDouble() == 1.0)
+                            finished++;
+                    }
+                }
+                sb.append(" (" + finished + " of " + size + " finished)");
+            }
+            statusPerHost.add(sb.toString());
+        }
+        return String.join("\n", statusPerHost);
+    }
+
+    @Command(name = "vespa-status-filedistribution", description = "Tool for getting file distribution status.")
+    public static class CommandLineArguments {
+
+        static CommandLineArguments build(String[] args) {
+            CommandLineArguments arguments = null;
+            try {
+                arguments = SingleCommand.singleCommand(CommandLineArguments.class).parse(args);
+            } catch (Exception e) {
+                System.err.println(e.getMessage());
+                System.err.println("Use --help to show usage.\n");
+                System.exit(1);
+            }
+
+            if (arguments.helpOption.showHelpIfRequested()) {
+                System.exit(0);
+            }
+
+            if (arguments.getTenantName() == null) {
+                System.err.println("'--tenant' not set.");
+                System.exit(1);
+            }
+
+            if (arguments.getApplicationName() == null) {
+                System.err.println("'--application' not set.");
+                System.exit(1);
+            }
+
+            return arguments;
+        }
+
+        @Inject
+        HelpOption helpOption;
+
+        @Option(name = {"--tenant"},
+                description = "tenant name")
+        private String tenantNameArg;
+
+        @Option(name = {"--application"},
+                description = "application name")
+        private String applicationNameArg;
+
+        @Option(name = {"--instance"},
+                description = "instance name")
+        private String instanceNameArg = "default";
+
+        @Option(name = {"--environment"},
+                description = "environment name")
+        private String environmentArg = "prod";
+
+        @Option(name = {"--region"},
+                description = "region name")
+        private String regionArg = "default";
+
+        @Option(name = {"--timeout"},
+                description = "The timeout (in seconds).")
+        private double timeoutArg = 5;
+
+        @Option(name = {"--debug"},
+                description = "Print debug log.")
+        private boolean debugArg;
+
+        public String getTenantName() { return tenantNameArg; }
+
+        public String getApplicationName() { return applicationNameArg; }
+
+        public String getInstanceName() { return instanceNameArg; }
+
+        public String getEnvironment() { return environmentArg; }
+
+        public String getRegion() { return regionArg; }
+
+        public double getTimeout() { return timeoutArg; }
+
+        public boolean getDebugFlag() { return debugArg; }
+    }
+
+}
diff --git a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClientTest.java b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClientTest.java
new file mode 100644
index 00000000000..fcbe880bfc7
--- /dev/null
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/status/FileDistributionStatusClientTest.java
@@ -0,0 +1,59 @@
+//  Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.vespa.filedistribution.status;
+
+import org.junit.Test;
+
+import static com.yahoo.vespa.filedistribution.status.FileDistributionStatusClient.CommandLineArguments;
+import static org.junit.Assert.assertEquals;
+
+public class FileDistributionStatusClientTest {
+
+    private static final CommandLineArguments arguments = createArguments("--tenant", "foo", "--application", "bar");
+    private final FileDistributionStatusClient client = new FileDistributionStatusClient(arguments);
+
+    @Test
+    public void finishedForAllHosts() {
+        String output = client.parseAndGenerateOutput("{\"status\":\"FINISHED\"}");
+        assertEquals("File distribution finished", output);
+    }
+
+    @Test
+    public void unknownForAllHosts() {
+        String output = client.parseAndGenerateOutput("{\"status\":\"UNKNOWN\", \"message\":\"Something went wrong\"}");
+        assertEquals("File distribution status unknown: Something went wrong", output);
+    }
+
+    @Test
+    public void manyHostsVariousStates() {
+        String statusForTwoHosts = createStatusForTwoHosts();
+        System.out.println(statusForTwoHosts);
+        String output = client.parseAndGenerateOutput(statusForTwoHosts);
+        assertEquals("File distribution in progress:\nlocalhost1: IN_PROGRESS (1 of 2 finished)\nlocalhost2: UNKNOWN (Connection timed out)", output);
+    }
+
+    private static CommandLineArguments createArguments(String... args) {
+        return CommandLineArguments.build(args);
+    }
+
+    private String createStatusForTwoHosts() {
+        return "{\"status\":\"IN_PROGRESS\"," +
+                "\"hosts\":[" + createInProgressStatusForHost("localhost1") + "," + createUnknownStatusForHost("localhost2") + "]" +
+                "}";
+    }
+
+    private String createInProgressStatusForHost(String hostname) {
+        return "{\"hostname\":\"" + hostname + "\"," +
+                "\"status\":\"IN_PROGRESS\"," +
+                "\"message\":\"\"," +
+                "\"fileReferences\":[" +
+                "{\"1234\":0.2}, {\"abcd\":1.0}]}";
+    }
+
+    private String createUnknownStatusForHost(String hostname) {
+        return "{\"hostname\":\"" + hostname + "\"," +
+                "\"status\":\"UNKNOWN\"," +
+                "\"message\":\"Connection timed out\"}";
+    }
+
+}