Prepare for multiple types of service dump artifacts

5 files changed, 114 insertions, 28 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ArtifactProducer.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ArtifactProducer.java
new file mode 100644
index 00000000000..24f070bbc35
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ArtifactProducer.java
@@ -0,0 +1,20 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.maintenance.servicedump;
+
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
+
+import java.io.IOException;
+
+/**
+ * Produces service dump artifacts.
+ *
+ * @author bjorncs
+ */
+interface ArtifactProducer {
+
+    String name();
+
+    void produceArtifact(NodeAgentContext context, String configId, UnixPath resultDirectoryInNode) throws IOException;
+
+}
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/JvmDumpProducer.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/JvmDumpProducer.java
new file mode 100644
index 00000000000..bc7703ba03a
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/JvmDumpProducer.java
@@ -0,0 +1,43 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.maintenance.servicedump;
+
+import com.yahoo.vespa.hosted.node.admin.container.ContainerOperations;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
+import com.yahoo.vespa.hosted.node.admin.task.util.process.CommandResult;
+
+import java.io.IOException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * Creates a dump of JVM based Vespa services using vespa-jvm-dumper
+ *
+ * @author bjorncs
+ */
+class JvmDumpProducer implements ArtifactProducer {
+
+    private static final Logger log = Logger.getLogger(JvmDumpProducer.class.getName());
+
+    private final ContainerOperations container;
+
+    JvmDumpProducer(ContainerOperations container) { this.container = container; }
+
+    public static String NAME = "jvm-dump";
+
+    @Override public String name() { return NAME; }
+
+    @Override
+    public void produceArtifact(NodeAgentContext context, String configId, UnixPath resultDirectoryInNode) throws IOException {
+        UnixPath vespaJvmDumper = new UnixPath(context.pathInNodeUnderVespaHome("bin/vespa-jvm-dumper"));
+        context.log(log, Level.INFO,
+                "Executing '" + vespaJvmDumper + "' with arguments '" + configId + "' and '" + resultDirectoryInNode + "'");
+        CommandResult result = container.executeCommandInContainerAsRoot(
+                context, vespaJvmDumper.toString(), configId, resultDirectoryInNode.toString());
+        context.log(log, Level.INFO,
+                "vespa-jvm-dumper exited with code '" + result.getExitCode() + "' and output:\n" + result.getOutput());
+        if (result.getExitCode() > 0) {
+            throw new IOException("Failed to jvm dump: " + result.getOutput());
+        }
+    }
+}
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ServiceDumpReport.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ServiceDumpReport.java
index 0134254b0c6..6ff4929ada1 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ServiceDumpReport.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/ServiceDumpReport.java
@@ -9,6 +9,8 @@ import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.reports.BaseReport;
 
+import java.util.List;
+
 /**
  * JSON representation of Vespa service dump report.
  *
@@ -27,6 +29,7 @@ class ServiceDumpReport extends BaseReport {
     private static final String CONFIG_ID_FIELD = "configId";
     private static final String EXPIRE_AT_FIELD = "expireAt";
     private static final String ERROR_FIELD = "error";
+    private static final String ARTIFACTS_FIELD = "artifacts";
 
     private final Long startedAt;
     private final Long completedAt;
@@ -35,6 +38,7 @@ class ServiceDumpReport extends BaseReport {
     private final String configId;
     private final Long expireAt;
     private final String error;
+    private final List<String> artifacts;
 
     @JsonCreator
     public ServiceDumpReport(@JsonProperty(CREATED_FIELD) Long createdAt,
@@ -44,7 +48,8 @@ class ServiceDumpReport extends BaseReport {
                              @JsonProperty(LOCATION_FIELD) String location,
                              @JsonProperty(CONFIG_ID_FIELD) String configId,
                              @JsonProperty(EXPIRE_AT_FIELD) Long expireAt,
-                             @JsonProperty(ERROR_FIELD) String error) {
+                             @JsonProperty(ERROR_FIELD) String error,
+                             @JsonProperty(ARTIFACTS_FIELD) List<String> artifacts) {
         super(createdAt, null);
         this.startedAt = startedAt;
         this.completedAt = completedAt;
@@ -53,6 +58,7 @@ class ServiceDumpReport extends BaseReport {
         this.configId = configId;
         this.expireAt = expireAt;
         this.error = error;
+        this.artifacts = artifacts;
     }
 
     @JsonGetter(STARTED_AT_FIELD) public Long startedAt() { return startedAt; }
@@ -62,6 +68,7 @@ class ServiceDumpReport extends BaseReport {
     @JsonGetter(CONFIG_ID_FIELD) public String configId() { return configId; }
     @JsonGetter(EXPIRE_AT_FIELD) public Long expireAt() { return expireAt; }
     @JsonGetter(ERROR_FIELD) public String error() { return error; }
+    @JsonGetter(ARTIFACTS_FIELD) public List<String> artifacts() { return artifacts; }
 
     @JsonIgnore public boolean isCompletedOrFailed() { return !isNullTimestamp(failedAt) || !isNullTimestamp(completedAt);  }
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImpl.java
index a0232e11d41..f98b47fa604 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImpl.java
@@ -11,14 +11,15 @@ import com.yahoo.vespa.hosted.node.admin.maintenance.sync.SyncFileInfo;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.FileFinder;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
-import com.yahoo.vespa.hosted.node.admin.task.util.process.CommandResult;
 
 import java.net.URI;
-import java.nio.file.Path;
 import java.time.Clock;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -38,6 +39,7 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
     private final SyncClient syncClient;
     private final NodeRepository nodeRepository;
     private final Clock clock;
+    private final Map<String, ArtifactProducer> artifactProducers;
 
     public VespaServiceDumperImpl(ContainerOperations container, SyncClient syncClient, NodeRepository nodeRepository) {
         this(container, syncClient, nodeRepository, Clock.systemUTC());
@@ -50,6 +52,9 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
         this.syncClient = syncClient;
         this.nodeRepository = nodeRepository;
         this.clock = clock;
+        this.artifactProducers = List.of(new JvmDumpProducer(container))
+                .stream()
+                .collect(Collectors.toMap(ArtifactProducer::name, Function.identity()));
     }
 
     @Override
@@ -76,11 +81,17 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
             handleFailure(context, request, startedAt, "Request already expired");
             return;
         }
+        List<String> artifactTypes = request.artifacts();
+        if (artifactTypes == null || artifactTypes.isEmpty()) {
+            handleFailure(context, request, startedAt, "No artifacts requested");
+            return;
+        }
         UnixPath directoryInNode = new UnixPath(context.pathInNodeUnderVespaHome("tmp/vespa-service-dump"));
         UnixPath directoryOnHost = new UnixPath(context.pathOnHostFromPathInNode(directoryInNode.toPath()));
         try {
             context.log(log, Level.INFO,
-                    "Creating dump for " + configId + " requested at " + Instant.ofEpochMilli(request.getCreatedMillisOrNull()));
+                    "Creating service dump for " + configId + " requested at "
+                            + Instant.ofEpochMilli(request.getCreatedMillisOrNull()));
             storeReport(context, createStartedReport(request, startedAt));
             if (directoryOnHost.exists()) {
                 context.log(log, Level.INFO, "Removing existing directory '" + directoryOnHost +"'.");
@@ -89,18 +100,23 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
             context.log(log, Level.INFO, "Creating '" + directoryOnHost +"'.");
             directoryOnHost.createDirectory();
             directoryOnHost.setPermissions("rwxrwxrwx");
-            UnixPath vespaJvmDumper = new UnixPath(context.pathInNodeUnderVespaHome("bin/vespa-jvm-dumper"));
-            context.log(log, Level.INFO, "Executing '" + vespaJvmDumper + "' with arguments '" + configId + "' and '" + directoryInNode + "'");
-            CommandResult result = container.executeCommandInContainerAsRoot(
-                    context, vespaJvmDumper.toString(), configId, directoryInNode.toString());
-            context.log(log, Level.INFO, "vespa-jvm-dumper exited with code '" + result.getExitCode() + "' and output:\n" + result.getOutput());
-            if (result.getExitCode() > 0) {
-                handleFailure(context, request, startedAt, "Failed to create dump: " + result.getOutput());
-                return;
-            }
+            List<SyncFileInfo> files = new ArrayList<>();
             URI destination = serviceDumpDestination(nodeSpec, createDumpId(request));
+            for (String artifactType : artifactTypes) {
+                ArtifactProducer producer = artifactProducers.get(artifactType);
+                if (producer == null) {
+                    handleFailure(context, request, startedAt, "No artifact producer exists for '" + artifactType + "'");
+                    return;
+                }
+                context.log(log, "Producing artifact of type '" + artifactType + "'");
+                UnixPath producerDirectoryOnHost = directoryOnHost.resolve(artifactType);
+                producerDirectoryOnHost.createDirectory();
+                producerDirectoryOnHost.setPermissions("rwxrwxrwx");
+                UnixPath producerDirectoryInNode = directoryInNode.resolve(artifactType);
+                producer.produceArtifact(context, configId, producerDirectoryInNode);
+                collectArtifactFilesToUpload(files, producerDirectoryOnHost, destination.resolve(artifactType + '/'), expiry);
+            }
             context.log(log, Level.INFO, "Uploading files with destination " + destination + " and expiry " + expiry);
-            List<SyncFileInfo> files = dumpFiles(directoryOnHost.toPath(), destination, expiry);
             if (!syncClient.sync(context, files, Integer.MAX_VALUE)) {
                 handleFailure(context, request, startedAt, "Unable to upload all files");
                 return;
@@ -117,10 +133,10 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
         }
     }
 
-    private List<SyncFileInfo> dumpFiles(Path directoryOnHost, URI destination, Instant expiry) {
-        return FileFinder.files(directoryOnHost).stream()
+    private void collectArtifactFilesToUpload(List<SyncFileInfo> files, UnixPath directoryOnHost, URI destination, Instant expiry) {
+        FileFinder.files(directoryOnHost.toPath()).stream()
                 .flatMap(file -> SyncFileInfo.forServiceDump(destination, file.path(), expiry).stream())
-                .collect(Collectors.toList());
+                .forEach(files::add);
     }
 
     private static Instant expireAt(Instant startedAt, ServiceDumpReport request) {
@@ -150,21 +166,21 @@ public class VespaServiceDumperImpl implements VespaServiceDumper {
     private static ServiceDumpReport createStartedReport(ServiceDumpReport request, Instant startedAt) {
         return new ServiceDumpReport(
                 request.getCreatedMillisOrNull(), startedAt.toEpochMilli(), null, null, null, request.configId(),
-                request.expireAt(), null);
+                request.expireAt(), null, request.artifacts());
     }
 
     private static ServiceDumpReport createSuccessReport(
             Clock clock, ServiceDumpReport request, Instant startedAt, URI location) {
         return new ServiceDumpReport(
                 request.getCreatedMillisOrNull(), startedAt.toEpochMilli(), clock.instant().toEpochMilli(), null,
-                location.toString(), request.configId(), request.expireAt(), null);
+                location.toString(), request.configId(), request.expireAt(), null, request.artifacts());
     }
 
     private static ServiceDumpReport createErrorReport(
             Clock clock, ServiceDumpReport request, Instant startedAt, String message) {
         return new ServiceDumpReport(
                 request.getCreatedMillisOrNull(), startedAt.toEpochMilli(), null, clock.instant().toEpochMilli(), null,
-                request.configId(), request.expireAt(), message);
+                request.configId(), request.expireAt(), message, request.artifacts());
     }
 
     static String createDumpId(ServiceDumpReport request) {
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
index 397305eae70..53e283abb42 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
@@ -41,7 +41,7 @@ class VespaServiceDumperImplTest {
     void creates_valid_dump_id_from_dump_request() {
         long nowMillis = Instant.now().toEpochMilli();
         ServiceDumpReport request = new ServiceDumpReport(
-                nowMillis, null, null, null, null, "default/container.3", null, null);
+                nowMillis, null, null, null, null, "default/container.3", null, null, List.of(JvmDumpProducer.NAME));
         String dumpId = VespaServiceDumperImpl.createDumpId(request);
         assertEquals("default-container-3-" + nowMillis, dumpId);
     }
@@ -59,7 +59,7 @@ class VespaServiceDumperImplTest {
         NodeRepoMock nodeRepository = new NodeRepoMock();
         ManualClock clock = new ManualClock(Instant.ofEpochMilli(1600001000000L));
         ServiceDumpReport request = new ServiceDumpReport(
-                1600000000000L, null, null, null, null, "default/container.1", null, null);
+                1600000000000L, null, null, null, null, "default/container.1", null, null, List.of(JvmDumpProducer.NAME));
         NodeSpec initialSpec = NodeSpec.Builder
                 .testSpec(HOSTNAME, NodeState.active)
                 .report(ServiceDumpReport.REPORT_ID, request.toJsonNode())
@@ -78,13 +78,13 @@ class VespaServiceDumperImplTest {
         String expectedJson =
                 "{\"createdMillis\":1600000000000,\"startedAt\":1600001000000,\"completedAt\":1600001000000," +
                         "\"location\":\"s3://uri-1/tenant1/service-dump/default-container-1-1600000000000/\"," +
-                        "\"configId\":\"default/container.1\"}";
+                        "\"configId\":\"default/container.1\",\"artifacts\":[\"jvm-dump\"]}";
         assertReportEquals(nodeRepository, expectedJson);
         verify(operations).executeCommandInContainerAsRoot(
-                context, "/opt/vespa/bin/vespa-jvm-dumper", "default/container.1", "/opt/vespa/tmp/vespa-service-dump");
+                context, "/opt/vespa/bin/vespa-jvm-dumper", "default/container.1", "/opt/vespa/tmp/vespa-service-dump/jvm-dump");
         List<URI> expectedUris = List.of(
-                URI.create("s3://uri-1/tenant1/service-dump/default-container-1-1600000000000/heap.bin.zst"),
-                URI.create("s3://uri-1/tenant1/service-dump/default-container-1-1600000000000/jstack"));
+                URI.create("s3://uri-1/tenant1/service-dump/default-container-1-1600000000000/jvm-dump/heap.bin.zst"),
+                URI.create("s3://uri-1/tenant1/service-dump/default-container-1-1600000000000/jvm-dump/jstack"));
         assertSyncedFiles(context, syncClient, expectedUris);
 
     }
@@ -113,8 +113,8 @@ class VespaServiceDumperImplTest {
         when(operations.executeCommandInContainerAsRoot(any(), any()))
                 .thenAnswer(invocation -> {
                     // Create dummy files to simulate vespa-jvm-dumper
-                    Files.createFile(tmpDirectory.resolve("vespa-service-dump/heap.bin"));
-                    Files.createFile(tmpDirectory.resolve("vespa-service-dump/jstack"));
+                    Files.createFile(tmpDirectory.resolve("vespa-service-dump/" + JvmDumpProducer.NAME + "/heap.bin"));
+                    Files.createFile(tmpDirectory.resolve("vespa-service-dump/" + JvmDumpProducer.NAME + "/jstack"));
                     return new CommandResult(null, 0, "result");
                 });
         return operations;