From 43d12011c2bd8c7680e4c00623a497a70304a487 Mon Sep 17 00:00:00 2001
From: Ming Wang <miwan@redhat.com>
Date: Thu, 31 Aug 2023 15:00:18 -0400
Subject: [PATCH] update or stop

---
 .../agent/remote/RecordingsContext.java       | 138 ++++++++++--------
 1 file changed, 76 insertions(+), 62 deletions(-)

diff --git a/src/main/java/io/cryostat/agent/remote/RecordingsContext.java b/src/main/java/io/cryostat/agent/remote/RecordingsContext.java
index 220229f9..e171d338 100644
--- a/src/main/java/io/cryostat/agent/remote/RecordingsContext.java
+++ b/src/main/java/io/cryostat/agent/remote/RecordingsContext.java
@@ -20,10 +20,14 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.net.MalformedURLException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.Map.Entry;
 import java.util.function.Consumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -32,13 +36,14 @@
 import javax.inject.Inject;
 
 import org.openjdk.jmc.common.unit.IConstrainedMap;
-import org.openjdk.jmc.common.unit.ITypedQuantity;
 import org.openjdk.jmc.common.unit.QuantityConversionException;
 import org.openjdk.jmc.common.unit.UnitLookup;
 import org.openjdk.jmc.flightrecorder.configuration.events.EventOptionID;
 import org.openjdk.jmc.flightrecorder.configuration.recording.RecordingOptionsBuilder;
+import org.openjdk.jmc.rjmx.ConnectionException;
 import org.openjdk.jmc.rjmx.ServiceNotAvailableException;
 import org.openjdk.jmc.rjmx.services.jfr.IFlightRecorderService;
+import org.openjdk.jmc.rjmx.services.jfr.IRecordingDescriptor;
 
 import io.cryostat.agent.StringUtils;
 import io.cryostat.core.FlightRecorderException;
@@ -113,19 +118,8 @@ public void handle(HttpExchange exchange) throws IOException {
                     break;
                 case "PATCH":
                     id = extractId(exchange);
-                    String request = requestStopOrUpdate(exchange);
-                    if (request == "STOPPED") {
-                        if (id >= 0) {
-                            handleStop(exchange, id);
-                        } else {
-                            exchange.sendResponseHeaders(HttpStatus.SC_BAD_REQUEST, BODY_LENGTH_NONE);
-                        }
-                    } else if (request == "UPDATE") {
-                        if (id >= 0) {
-                            handleRecordingUpdate(exchange, id);
-                        } else {
-                            exchange.sendResponseHeaders(HttpStatus.SC_BAD_REQUEST, BODY_LENGTH_NONE);
-                        }
+                    if (id >= 0) {
+                        handleStopOrUpdate(exchange, id);
                     } else {
                         exchange.sendResponseHeaders(HttpStatus.SC_BAD_REQUEST, BODY_LENGTH_NONE);
                     }
@@ -149,14 +143,6 @@ public void handle(HttpExchange exchange) throws IOException {
         }
     }
 
-    private static String requestStopOrUpdate(HttpExchange exchange) throws IOException {
-        try (InputStream body = exchange.getRequestBody()) {
-            ObjectMapper mapper = new ObjectMapper();
-            JsonNode jsonMap = mapper.readTree(body);
-            return jsonMap.get("state").toString();
-        }
-    }
-
     private static long extractId(HttpExchange exchange) throws IOException {
         Matcher m = PATH_ID_PATTERN.matcher(exchange.getRequestURI().getPath());
         if (!m.find()) {
@@ -253,53 +239,81 @@ private void handleStartRecordingOrSnapshot(HttpExchange exchange) throws IOExce
         }
     }
 
-    private void handleRecordingUpdate(HttpExchange exchange, long id) throws IOException {
-        try (InputStream body = exchange.getRequestBody()) {
-            ObjectMapper mapper = new ObjectMapper();
+    private void handleStopOrUpdate(HttpExchange exchange, long id) throws IOException {
+        try {
+            JFRConnection conn =
+                    jfrConnectionToolkit.connect(
+                            jfrConnectionToolkit.createServiceURL("localhost", 0));
+            IFlightRecorderService svc = conn.getService();
+            IRecordingDescriptor dsc = svc.getAvailableRecordings().stream().filter(r -> r.getId() == id).findFirst().get();
+            RecordingOptionsBuilder builder = new RecordingOptionsBuilder(conn.getService());
+            
+            InputStream body = exchange.getRequestBody();
             JsonNode jsonMap = mapper.readTree(body);
-            invokeOnRecording(
-                    exchange,
-                    id,
-                    r -> {
-                        try {
-                            if (jsonMap.has("name")) {
-                                r.setName(jsonMap.get("name").toString());
-                            }
-                            if (jsonMap.has("toDisk")) {
-                                r.setToDisk(jsonMap.get("toDisk").asBoolean());
-                            }
-                            if (jsonMap.has("duration")) {
-                                r.setDuration(Duration.ofMillis(jsonMap.get("duration").asLong()));
-                            }
-                            if (jsonMap.has("maxSize")) {
-                                r.setMaxSize(jsonMap.get("name").asLong());
-                            }
-                            if (jsonMap.has("maxAge")) {
-                                r.setMaxAge(Duration.ofMillis(jsonMap.get("name").asLong()));
-                            }
-                        } catch (IllegalStateException e) {
-                            sendHeader(exchange, HttpStatus.SC_CONFLICT);
+            Iterator<Entry<String, JsonNode>> fields = jsonMap.fields();
+
+            while (fields.hasNext()) {
+                Map.Entry<String, JsonNode> field = fields.next();
+
+                switch(field.getKey()){
+                    case "state":
+                        if ("STOPPED".equals(field.getValue().toString())) {
+                            handleStop(exchange, id);
+                            break;
+                        } else if (!"UPDATE".equals(field.getValue().toString())) {
+                            exchange.sendResponseHeaders(HttpStatus.SC_BAD_REQUEST, BODY_LENGTH_NONE);
                         }
-                    });
+                        break;
+                    case "name":
+                        builder = builder.name(field.getValue().toString());
+                        break;
+                    case "duration":
+                        builder = builder.duration(field.getValue().asLong());
+                        break;
+                    case "maxSize":
+                        builder = builder.maxSize(field.getValue().asLong());
+                        break;
+                    case "maxAge":
+                        builder = builder.maxAge(field.getValue().asLong());
+                        break;
+                    case "toDisk":
+                        builder = builder.toDisk(field.getValue().asBoolean());
+                        break;
+                    default:
+                        log.warn("Unknown recording option {}", field.getKey());
+                        exchange.sendResponseHeaders(
+                                HttpStatus.SC_METHOD_NOT_ALLOWED, BODY_LENGTH_NONE);
+                        break;
+                }
+            }
+            svc.updateRecordingOptions(dsc, builder.build());
+            exchange.sendResponseHeaders(HttpStatus.SC_CREATED, BODY_LENGTH_UNKNOWN);
+        } catch ( ServiceNotAvailableException
+                | org.openjdk.jmc.rjmx.services.jfr.FlightRecorderException
+                | QuantityConversionException e){
+            log.error("Failed to update recording", e);
+            exchange.sendResponseHeaders(HttpStatus.SC_INTERNAL_SERVER_ERROR, BODY_LENGTH_NONE);
+        } finally {
+            exchange.close();
         }
     }
 
     private void handleStop(HttpExchange exchange, long id) throws IOException {
         invokeOnRecording(
-                exchange,
-                id,
-                r -> {
-                    try {
-                        boolean stopped = r.stop();
-                        if (!stopped) {
-                            sendHeader(exchange, HttpStatus.SC_BAD_REQUEST);
-                        } else {
-                            sendHeader(exchange, HttpStatus.SC_NO_CONTENT);
-                        }
-                    } catch (IllegalStateException e) {
-                        sendHeader(exchange, HttpStatus.SC_CONFLICT);
+            exchange,
+            id,
+            r -> {
+                try {
+                    boolean stopped = r.stop();
+                    if (!stopped) {
+                        sendHeader(exchange, HttpStatus.SC_BAD_REQUEST);
+                    } else {
+                        sendHeader(exchange, HttpStatus.SC_NO_CONTENT);
                     }
-                });
+                } catch (IllegalStateException e) {
+                    sendHeader(exchange, HttpStatus.SC_CONFLICT);
+                }
+            });
     }
 
     private void handleDelete(HttpExchange exchange, long id) throws IOException {
@@ -352,7 +366,7 @@ private SerializableRecordingDescriptor startRecording(StartRecordingRequest req
             throws QuantityConversionException, ServiceNotAvailableException,
                     FlightRecorderException,
                     org.openjdk.jmc.rjmx.services.jfr.FlightRecorderException,
-                    InvalidEventTemplateException, InvalidXmlException, IOException {
+                    InvalidEventTemplateException, InvalidXmlException, IOException, FlightRecorderException {
         Runnable cleanup = () -> {};
         try {
             JFRConnection conn =