From a9782e6b71760a4a5ba988ede4d01fbde7abac49 Mon Sep 17 00:00:00 2001
From: brig <ybrigo@gmail.com>
Date: Mon, 2 Oct 2023 16:21:54 +0200
Subject: [PATCH] runtime-v2: automatically convert nonserializable map.entry
 to serizalizable

---
 .../v2/runner/script/VariablesSanitizer.java  | 20 +++++++++++--------
 .../v2/runner/vm/ExpressionCommand.java       |  2 ++
 .../concord/runtime/v2/runner/MainTest.java   | 11 ++++++++++
 .../runner/entrySetSerialization/concord.yaml | 10 ++++++++++
 4 files changed, 35 insertions(+), 8 deletions(-)
 create mode 100644 runtime/v2/runner/src/test/resources/com/walmartlabs/concord/runtime/v2/runner/entrySetSerialization/concord.yaml

diff --git a/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/script/VariablesSanitizer.java b/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/script/VariablesSanitizer.java
index d635e75933..41a7fe7418 100644
--- a/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/script/VariablesSanitizer.java
+++ b/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/script/VariablesSanitizer.java
@@ -20,30 +20,34 @@
  * =====
  */
 
+import java.io.Serializable;
 import java.util.*;
 import java.util.stream.Collectors;
 
 public final class VariablesSanitizer {
 
     @SuppressWarnings("unchecked")
-    public static Object sanitize(Object scriptObj) {
-        if (scriptObj instanceof Set) {
-            Set<Object> c = (Set<Object>) scriptObj;
+    public static Object sanitize(Object obj) {
+        if (obj instanceof Set) {
+            Set<Object> c = (Set<Object>) obj;
             return c.stream()
                     .map(VariablesSanitizer::sanitize)
                     .collect(Collectors.toSet());
-        } else if (scriptObj instanceof Collection) {
-            Collection<Object> c = (Collection<Object>) scriptObj;
+        } else if (obj instanceof Collection) {
+            Collection<Object> c = (Collection<Object>) obj;
             return c.stream()
                     .map(VariablesSanitizer::sanitize)
                     .collect(Collectors.toList());
-        } else if (scriptObj instanceof Map) {
-            Map<Object, Object> m = (Map<Object, Object>) scriptObj;
+        } else if (obj instanceof Map) {
+            Map<Object, Object> m = (Map<Object, Object>) obj;
             Map<Object, Object> result = new LinkedHashMap<>();
             m.forEach((key, value) -> result.put(sanitize(key), sanitize(value)));
             return result;
+        } else if (obj instanceof Map.Entry && !(obj instanceof Serializable)) {
+            Map.Entry<?, ?> e = (Map.Entry<?, ?>) obj;
+            return new AbstractMap.SimpleEntry<>(e.getKey(), e.getValue());
         } else {
-            return scriptObj;
+            return obj;
         }
     }
 
diff --git a/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/vm/ExpressionCommand.java b/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/vm/ExpressionCommand.java
index 05c6529b70..fa8325eb67 100644
--- a/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/vm/ExpressionCommand.java
+++ b/runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/vm/ExpressionCommand.java
@@ -22,6 +22,7 @@
 
 import com.walmartlabs.concord.runtime.v2.model.Expression;
 import com.walmartlabs.concord.runtime.v2.model.ExpressionOptions;
+import com.walmartlabs.concord.runtime.v2.runner.script.VariablesSanitizer;
 import com.walmartlabs.concord.runtime.v2.sdk.EvalContextFactory;
 import com.walmartlabs.concord.runtime.v2.sdk.ExpressionEvaluator;
 import com.walmartlabs.concord.runtime.v2.sdk.Context;
@@ -58,6 +59,7 @@ protected void execute(Runtime runtime, State state, ThreadId threadId) {
         EvalContextFactory ecf = runtime.getService(EvalContextFactory.class);
         ExpressionEvaluator ee = runtime.getService(ExpressionEvaluator.class);
         Object result = ee.eval(ecf.global(ctx), expr, Object.class);
+        result = VariablesSanitizer.sanitize(result);
 
         ExpressionOptions opts = Objects.requireNonNull(step.getOptions());
         if (!opts.outExpr().isEmpty()) {
diff --git a/runtime/v2/runner/src/test/java/com/walmartlabs/concord/runtime/v2/runner/MainTest.java b/runtime/v2/runner/src/test/java/com/walmartlabs/concord/runtime/v2/runner/MainTest.java
index b804e05c10..e23b9ffa3e 100644
--- a/runtime/v2/runner/src/test/java/com/walmartlabs/concord/runtime/v2/runner/MainTest.java
+++ b/runtime/v2/runner/src/test/java/com/walmartlabs/concord/runtime/v2/runner/MainTest.java
@@ -1530,6 +1530,17 @@ public void testArrayEvalSerialize() throws Exception {
         assertLog(log, ".*" + Pattern.quote("{dev=dev-cloud1}, {prod=prod-cloud1}, {test=test-cloud1}, {perf=perf-cloud2}, {ci=perf-ci}") + ".*");
     }
 
+    @Test
+    public void testEntrySetSerialization() throws Exception {
+        deploy("entrySetSerialization");
+
+        save(ProcessConfiguration.builder()
+                .build());
+
+        byte[] log = run();
+        assertLog(log, ".*myList: \\[k=v\\].*");
+    }
+
     private void deploy(String resource) throws URISyntaxException, IOException {
         Path src = Paths.get(MainTest.class.getResource(resource).toURI());
         IOUtils.copy(src, workDir);
diff --git a/runtime/v2/runner/src/test/resources/com/walmartlabs/concord/runtime/v2/runner/entrySetSerialization/concord.yaml b/runtime/v2/runner/src/test/resources/com/walmartlabs/concord/runtime/v2/runner/entrySetSerialization/concord.yaml
new file mode 100644
index 0000000000..0a4e3192b7
--- /dev/null
+++ b/runtime/v2/runner/src/test/resources/com/walmartlabs/concord/runtime/v2/runner/entrySetSerialization/concord.yaml
@@ -0,0 +1,10 @@
+flows:
+  default:
+    - set:
+        myMap:
+          k: v
+
+    - expr: ${myMap.entrySet().stream().toList()}
+      out: myList
+
+    - log: "myList: ${myList}"
\ No newline at end of file