fixup! feat(recordings): set recording template in label

cryostatio · Mar 16, 2022 · 68267fc · 68267fc
1 parent 3f60c43
commit 68267fc
Showing 1 changed file with 39 additions and 34 deletions.
diff --git a/src/main/java/io/cryostat/recordings/RecordingTargetHelper.java b/src/main/java/io/cryostat/recordings/RecordingTargetHelper.java
@@ -63,6 +63,7 @@
 
 import io.cryostat.core.log.Logger;
 import io.cryostat.core.net.JFRConnection;
+import io.cryostat.core.templates.Template;
 import io.cryostat.core.templates.TemplateType;
 import io.cryostat.jmc.serialization.HyperlinkedSerializableRecordingDescriptor;
 import io.cryostat.messaging.notifications.NotificationFactory;
@@ -139,6 +140,8 @@ public IRecordingDescriptor startRecording(
         return targetConnectionManager.executeConnectedTask(
                 connectionDescriptor,
                 connection -> {
+                    TemplateType preferredTemplateType =
+                            getPreferredTemplateType(connection, templateName, templateType);
                     Optional<IRecordingDescriptor> previous =
                             getDescriptorByName(connection, recordingName);
                     if (previous.isPresent()) {
@@ -156,7 +159,10 @@ public IRecordingDescriptor startRecording(
                                     .getService()
                                     .start(
                                             recordingOptions,
-                                            enableEvents(connection, templateName, templateType));
+                                            enableEvents(
+                                                    connection,
+                                                    templateName,
+                                                    preferredTemplateType));
                     String targetId = connectionDescriptor.getTargetId();
                     Metadata metadata =
                             recordingMetadataManager.getMetadata(targetId, recordingName);
@@ -165,9 +171,7 @@ public IRecordingDescriptor startRecording(
                             "template",
                             String.format(
                                     "template=%s,type=%s",
-                                    templateName,
-                                    (templateType == null ? TemplateType.TARGET : templateType)
-                                            .name()));
+                                    templateName, preferredTemplateType.name()));
                     metadata = new Metadata(labels);
                     metadata =
                             recordingMetadataManager
@@ -461,45 +465,46 @@ private void cancelScheduledNotificationIfExists(String targetId, String stopped
         }
     }
 
+    private TemplateType getPreferredTemplateType(
+            JFRConnection connection, String templateName, TemplateType templateType)
+            throws Exception {
+        if (templateType != null) {
+            return templateType;
+        }
+        if (templateName.equals("ALL")) {
+            // special case for the ALL meta-template
+            return TemplateType.TARGET;
+        }
+        List<Template> templates = connection.getTemplateService().getTemplates();
+        boolean custom =
+                templates.stream()
+                        .filter(t -> t.getName().equals(templateName))
+                        .anyMatch(t -> t.getType().equals(TemplateType.CUSTOM));
+        if (custom) {
+            return TemplateType.CUSTOM;
+        }
+        boolean target =
+                templates.stream()
+                        .filter(t -> t.getName().equals(templateName))
+                        .anyMatch(t -> t.getType().equals(TemplateType.TARGET));
+        if (target) {
+            return TemplateType.TARGET;
+        }
+        throw new IllegalArgumentException(
+                String.format("Invalid/unknown event template %s", templateName));
+    }
+
     private IConstrainedMap<EventOptionID> enableEvents(
             JFRConnection connection, String templateName, TemplateType templateType)
             throws Exception {
         if (templateName.equals("ALL")) {
             return enableAllEvents(connection);
         }
-        if (templateType != null) {
-            return connection
-                    .getTemplateService()
-                    .getEvents(templateName, templateType)
-                    .orElseThrow(
-                            () ->
-                                    new IllegalArgumentException(
-                                            String.format(
-                                                    "No template \"%s\" found with type %s",
-                                                    templateName, templateType)));
-        }
         // if template type not specified, try to find a Custom template by that name. If none,
         // fall back on finding a Target built-in template by the name. If not, throw an
         // exception and bail out.
-        return connection
-                .getTemplateService()
-                .getEvents(templateName, TemplateType.CUSTOM)
-                .or(
-                        () -> {
-                            try {
-                                return connection
-                                        .getTemplateService()
-                                        .getEvents(templateName, TemplateType.TARGET);
-                            } catch (Exception e) {
-                                return Optional.empty();
-                            }
-                        })
-                .orElseThrow(
-                        () ->
-                                new IllegalArgumentException(
-                                        String.format(
-                                                "Invalid/unknown event template %s",
-                                                templateName)));
+        TemplateType type = getPreferredTemplateType(connection, templateName, templateType);
+        return connection.getTemplateService().getEvents(templateName, type).get();
     }
 
     private IConstrainedMap<EventOptionID> enableAllEvents(JFRConnection connection)