fix(recordings): correctly retrieve custom event templates from S3 st…

…orage (#492)
cryostatio · Jun 3, 2024 · 65f9b30 · 65f9b30
1 parent 4546260
commit 65f9b30
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 30 deletions.
diff --git a/src/main/java/io/cryostat/recordings/RecordingHelper.java b/src/main/java/io/cryostat/recordings/RecordingHelper.java
@@ -369,11 +369,22 @@ private Uni<ActiveRecording> startRecordingImpl(
                             }
                             IConstrainedMap<String> recordingOptions = optionsBuilder.build();
 
-                            return conn.getService()
-                                    .start(
-                                            recordingOptions,
-                                            template.getName(),
-                                            template.getType());
+                            switch (template.getType()) {
+                                case CUSTOM:
+                                    return conn.getService()
+                                            .start(
+                                                    recordingOptions,
+                                                    customTemplateService
+                                                            .getXml(
+                                                                    template.getName(),
+                                                                    TemplateType.CUSTOM)
+                                                            .orElseThrow());
+                                case TARGET:
+                                    return conn.getService().start(recordingOptions, template);
+                                default:
+                                    throw new IllegalStateException(
+                                            "Unknown template type: " + template.getType());
+                            }
                         });
 
         Map<String, String> labels = new HashMap<>(rawLabels);

diff --git a/src/main/java/io/cryostat/targets/AgentConnection.java b/src/main/java/io/cryostat/targets/AgentConnection.java
@@ -36,6 +36,7 @@
 import io.cryostat.core.sys.Clock;
 import io.cryostat.core.templates.RemoteTemplateService;
 import io.cryostat.core.templates.TemplateService;
+import io.cryostat.events.S3TemplateService;
 
 import io.smallrye.common.annotation.Blocking;
 import jakarta.enterprise.context.ApplicationScoped;
@@ -45,10 +46,12 @@
 class AgentConnection implements JFRConnection {
 
     private final AgentClient client;
+    private final TemplateService customTemplateService;
     private final Logger logger = Logger.getLogger(getClass());
 
-    AgentConnection(AgentClient client) {
+    AgentConnection(AgentClient client, TemplateService customTemplateService) {
         this.client = client;
+        this.customTemplateService = customTemplateService;
     }
 
     @Override
@@ -113,7 +116,7 @@ public int getPort() {
     @Override
     public CryostatFlightRecorderService getService()
             throws ConnectionException, IOException, ServiceNotAvailableException {
-        return new AgentJFRService(client, getTemplateService());
+        return new AgentJFRService(client, customTemplateService);
     }
 
     @Override
@@ -141,10 +144,11 @@ public MBeanMetrics getMBeanMetrics()
     public static class Factory {
 
         @Inject AgentClient.Factory clientFactory;
+        @Inject S3TemplateService customTemplateService;
         @Inject Logger logger;
 
         public AgentConnection createConnection(Target target) {
-            return new AgentConnection(clientFactory.create(target));
+            return new AgentConnection(clientFactory.create(target), customTemplateService);
         }
     }
 }
diff --git a/src/main/java/io/cryostat/targets/AgentJFRService.java b/src/main/java/io/cryostat/targets/AgentJFRService.java
@@ -19,6 +19,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.text.ParseException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -46,6 +47,7 @@
 import io.cryostat.core.EventOptionsBuilder.EventOptionException;
 import io.cryostat.core.EventOptionsBuilder.EventTypeException;
 import io.cryostat.core.net.CryostatFlightRecorderService;
+import io.cryostat.core.templates.Template;
 import io.cryostat.core.templates.TemplateService;
 import io.cryostat.core.templates.TemplateType;
 
@@ -217,21 +219,61 @@ public void updateRecordingOptions(
 
     @Blocking
     @Override
-    public IRecordingDescriptor start(
-            IConstrainedMap<String> recordingOptions,
-            String templateName,
-            TemplateType preferredTemplateType)
+    public IRecordingDescriptor start(IConstrainedMap<String> recordingOptions, String template)
+            throws FlightRecorderException,
+                    ParseException,
+                    IOException,
+                    QuantityConversionException {
+        long duration =
+                (Optional.ofNullable(
+                                        (ITypedQuantity)
+                                                recordingOptions.get(
+                                                        RecordingOptionsBuilder.KEY_DURATION))
+                                .orElse(UnitLookup.MILLISECOND.quantity(0)))
+                        .longValueIn(UnitLookup.MILLISECOND);
+        long maxSize =
+                (Optional.ofNullable(
+                                        (ITypedQuantity)
+                                                recordingOptions.get(
+                                                        RecordingOptionsBuilder.KEY_MAX_SIZE))
+                                .orElse(UnitLookup.BYTE.quantity(0)))
+                        .longValueIn(UnitLookup.BYTE);
+        long maxAge =
+                (Optional.ofNullable(
+                                        (ITypedQuantity)
+                                                recordingOptions.get(
+                                                        RecordingOptionsBuilder.KEY_MAX_AGE))
+                                .orElse(UnitLookup.MILLISECOND.quantity(0)))
+                        .longValueIn(UnitLookup.MILLISECOND);
+        StartRecordingRequest req =
+                new StartRecordingRequest(
+                        recordingOptions.get("name").toString(),
+                        null,
+                        template,
+                        duration,
+                        maxSize,
+                        maxAge);
+        return client.startRecording(req).await().atMost(client.getTimeout());
+    }
+
+    @Blocking
+    @Override
+    public IRecordingDescriptor start(IConstrainedMap<String> recordingOptions, Template template)
             throws io.cryostat.core.FlightRecorderException,
                     FlightRecorderException,
                     ConnectionException,
+                    ParseException,
                     IOException,
                     FlightRecorderException,
                     ServiceNotAvailableException,
                     QuantityConversionException,
                     EventOptionException,
                     EventTypeException {
-        StartRecordingRequest req;
-        String recordingName = recordingOptions.get("name").toString();
+        if (template.getType().equals(TemplateType.CUSTOM)) {
+            return start(
+                    recordingOptions,
+                    templateService.getXml(template.getName(), template.getType()).orElseThrow());
+        }
         long duration =
                 (Optional.ofNullable(
                                         (ITypedQuantity)
@@ -253,22 +295,14 @@ public IRecordingDescriptor start(
                                                         RecordingOptionsBuilder.KEY_MAX_AGE))
                                 .orElse(UnitLookup.MILLISECOND.quantity(0)))
                         .longValueIn(UnitLookup.MILLISECOND);
-        if (preferredTemplateType.equals(TemplateType.CUSTOM)) {
-            req =
-                    new StartRecordingRequest(
-                            recordingName,
-                            null,
-                            templateService
-                                    .getXml(templateName, preferredTemplateType)
-                                    .orElseThrow(),
-                            duration,
-                            maxSize,
-                            maxAge);
-        } else {
-            req =
-                    new StartRecordingRequest(
-                            recordingName, templateName, null, duration, maxSize, maxAge);
-        }
+        StartRecordingRequest req =
+                new StartRecordingRequest(
+                        recordingOptions.get("name").toString(),
+                        template.getName(),
+                        null,
+                        duration,
+                        maxSize,
+                        maxAge);
         return client.startRecording(req).await().atMost(client.getTimeout());
     }